You are viewing a plain text version of this content. The canonical link for it is here.
Posted to pr@cassandra.apache.org by "mike-tr-adamson (via GitHub)" <gi...@apache.org> on 2023/09/06 16:11:41 UTC
[GitHub] [cassandra] mike-tr-adamson opened a new pull request, #2673: Add support for vector search in SAI
mike-tr-adamson opened a new pull request, #2673:
URL: https://github.com/apache/cassandra/pull/2673
Thanks for sending a pull request! Here are some tips if you're new here:
* Ensure you have added or run the [appropriate tests](https://cassandra.apache.org/_/development/testing.html) for your PR.
* Be sure to keep the PR description updated to reflect all changes.
* Write your PR title to summarize what this PR proposes.
* If possible, provide a concise example to reproduce the issue for a faster review.
* Read our [contributor guidelines](https://cassandra.apache.org/_/development/index.html)
* If you're making a documentation change, see our [guide to documentation contribution](https://cassandra.apache.org/_/development/documentation.html)
Commit messages should follow the following format:
```
<One sentence description, usually Jira title or CHANGES.txt summary>
<Optional lengthier description (context on patch)>
patch by <Authors>; reviewed by <Reviewers> for CASSANDRA-#####
Co-authored-by: Name1 <email1>
Co-authored-by: Name2 <email2>
```
The [Cassandra Jira](https://issues.apache.org/jira/projects/CASSANDRA/issues/)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1334280798
##########
src/java/org/apache/cassandra/cql3/Operator.java:
##########
@@ -258,6 +258,20 @@ public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteB
{
throw new UnsupportedOperationException();
}
+ },
+ ANN(15)
+ {
+ @Override
+ public String toString()
+ {
+ return "ANN";
+ }
+
+ @Override
+ public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteBuffer rightOperand)
+ {
+ return true;
Review Comment:
I have changed the implementation to throw a UOE and documented the reason. This method should never be called because the ANN operator is only supported by the index. This in the same as the LIKE operator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1336152655
##########
src/java/org/apache/cassandra/cql3/SingleColumnRelation.java:
##########
@@ -264,6 +265,16 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
return new SingleColumnRestriction.LikeRestriction(columnDef, operator, term);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ ColumnMetadata columnDef = table.getExistingColumn(entity);
+ if (!(columnDef.type instanceof VectorType))
+ throw invalidRequest("ANN is only supported against DENSE FLOAT32 columns");
Review Comment:
Just a note: I have added extra unit tests to `VectorInvalidQueryTest` for the cases that were missing. These missing tests highlighted the fact that this code was no longer needed. I have refactored the `StatementRestrictions` initialization where it relates to vector indexes and ANN ordering to correctly handle the missing test cases.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341262017
##########
src/java/org/apache/cassandra/db/marshal/VectorType.java:
##########
@@ -365,6 +366,7 @@ public abstract class VectorSerializer extends TypeSerializer<List<T>>
public abstract <V> List<V> split(V buffer, ValueAccessor<V> accessor);
public abstract <V> V serializeRaw(List<V> elements, ValueAccessor<V> accessor);
+ public abstract float[] deserializeFloatArray(ByteBuffer input);
Review Comment:
No, `composeAsFloat` composes a `float[]` for the specific `VectorType`. As in, it uses the `dimension` specified for the `VectorType` instance this is called on. The `deserializeFloatArray` method is not specific to single instance and gets the dimension from the buffer passed to is that was previously generated by a call to `serialize`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341243973
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1133,13 +1149,13 @@ public SelectStatement prepare(ClientState state, boolean forView) throws Invali
List<Selectable> selectables = RawSelector.toSelectables(selectClause, table);
boolean containsOnlyStaticColumns = selectOnlyStaticColumns(table, selectables);
- StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, containsOnlyStaticColumns, forView);
+ List<Ordering> orderings = getOrderings(table);
+ StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, orderings, containsOnlyStaticColumns, forView);
// If we order post-query, the sorted column needs to be in the ResultSet for sorting,
// even if we don't ultimately ship them to the client (CASSANDRA-4911).
- Map<ColumnMetadata, Boolean> orderingColumns = getOrderingColumns(table);
- Set<ColumnMetadata> resultSetOrderingColumns = restrictions.keyIsInRelation() ? orderingColumns.keySet()
- : Collections.emptySet();
+ Map<ColumnMetadata, Ordering> orderingColumns = getOrderingColumns(orderings);
+ Set<ColumnMetadata> resultSetOrderingColumns = getResultSetOrdering(restrictions, orderingColumns);
Review Comment:
This is a bug in the distributed test code. Specifically, `RowUtil.toQueryResult`. This is incorrectly using `rows.result.metadata.names` instead of `rows.result.metadata.requestNames()`. The `requestNames` method specifically removes any columns that have not been requested. I have raised CASSANDRA-18892 and attached a fix patch there for this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327395479
##########
src/java/org/apache/cassandra/index/sai/VectorContext.java:
##########
@@ -0,0 +1,168 @@
+/*
+ * 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.index.sai;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnHeapHnsw;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.lucene.util.Bits;
+
+public class VectorContext
+{
+ public int hnswVectorsAccessed;
+ public int hnswVectorCacheHits;
Review Comment:
Can we add some JavaDoc for these attributes?
##########
src/java/org/apache/cassandra/index/Index.java:
##########
@@ -440,6 +445,19 @@ default boolean filtersMultipleContains()
*/
public RowFilter getPostIndexQueryFilter(RowFilter filter);
+ /**
+ * Return a comparator that reorders query result before sending to client
+ *
+ * @param restriction restriction that requires current index
+ * @param columnIndex idx of the indexed column in returned row
+ * @param options query options
+ * @return comparator that for post-query ordering; or null if not supported
+ */
+ default Comparator<List<ByteBuffer>> getPostQueryOrdering(Restriction restriction, int columnIndex, QueryOptions options)
Review Comment:
Nit: add `@Nullable`
##########
src/java/org/apache/cassandra/index/sai/VectorContext.java:
##########
@@ -0,0 +1,168 @@
+/*
+ * 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.index.sai;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnHeapHnsw;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.lucene.util.Bits;
+
+public class VectorContext
Review Comment:
Can we add some class-level javaDoc?
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java:
##########
@@ -0,0 +1,48 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.PostingList;
+
+/**
+ * There are two steps in ordering:
+ * <p>
+ * 1. Limit a single sstable's results to the correct keys. At this stage
+ * we put them back in primary key order to play nice with the rest of
+ * the query pipeline.
+ * 2. Merge the results from multiple sstables. Now we leave them in the
+ * final, correct order.
+ * <p>
+ * SegmentOrdering handles the first step.
+ */
+public interface SegmentOrdering
+{
+ /**
+ * Reorder, limit, and put back into original order the results from a single sstable
+ */
+ default KeyRangeIterator limitToTopResults(QueryContext context, PostingList iterator, Expression exp) throws IOException
Review Comment:
Maybe we could name this `limitToTopKResults`, to be consistent with the many `topK` vars and methods around?
##########
src/java/org/apache/cassandra/index/sai/disk/SSTableIndex.java:
##########
@@ -121,6 +123,24 @@ public SSTableIndex(SSTableContext sstableContext, IndexContext indexContext)
*/
public abstract AbstractBounds<PartitionPosition> bounds();
+// /**
+// * Perform a search on the index for a single expression and keyRange.
+// * <p>
+// * The result is a {@link List} of {@link KeyRangeIterator} because there will
+// * be a {@link KeyRangeIterator} for each segment in the index. The result
+// * will never be null but may be an empty {@link List}.
+// *
+// * @param expression The {@link Expression} to be searched for
+// * @param keyRange The {@link AbstractBounds<PartitionPosition>} defining the
+// * token range for the search
+// * @param context The {@link QueryContext} holding the per-query state
+// * @return a {@link List} of {@link KeyRangeIterator}s containing the results
+// * of the search
+// */
+// public abstract List<KeyRangeIterator<PrimaryKey>> search(Expression expression,
+// AbstractBounds<PartitionPosition> keyRange,
+// QueryContext context) throws IOException;
Review Comment:
This seems accidentally included
##########
src/java/org/apache/cassandra/index/sai/disk/RowMapping.java:
##########
@@ -163,6 +177,18 @@ public void add(PrimaryKey key, long sstableRowId) throws InMemoryTrie.SpaceExha
if (minKey == null)
minKey = key;
maxKey = key;
+ count++;
+ }
+
+ public int get(PrimaryKey key)
Review Comment:
Nit: I'd add JavaDoc saying that `-1` means not found.
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/Segment.java:
##########
@@ -97,6 +98,19 @@ public long indexFileCacheSize()
return index == null ? 0 : index.indexFileCacheSize();
}
+// /**
+// * Search on-disk index synchronously
+// *
+// * @param expression to filter on disk index
+// * @param context to track per sstable cache and per query metrics
+//
+// * @return range iterator that matches given expression
+// */
+// public KeyRangeIterator<PrimaryKey> search(Expression expression, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
+// {
+// return index.search(expression, keyRange, context);
+// }
Review Comment:
This seems accidentally included.
##########
src/java/org/apache/cassandra/index/sai/disk/v1/postings/ReorderingPostingList.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.index.sai.disk.v1.postings;
+
+import java.io.IOException;
+import java.util.PrimitiveIterator;
+
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.lucene.util.LongHeap;
+
+public class ReorderingPostingList implements PostingList
+{
+ private final LongHeap segmentRowIds;
+ private final int size;
+
+ public ReorderingPostingList(PrimitiveIterator.OfInt source, int estimatedSize) throws IOException
Review Comment:
Nit: doesn't throw `IOException`
##########
src/java/org/apache/cassandra/index/sai/disk/v1/IndexWriterConfig.java:
##########
@@ -0,0 +1,164 @@
+/*
+ * 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.index.sai.disk.v1;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.apache.cassandra.config.CassandraRelevantProperties;
+import org.apache.cassandra.db.marshal.AbstractType;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+import static org.apache.cassandra.config.CassandraRelevantProperties.SAI_VECTOR_SEARCH_MAX_TOP_K;
+
+/**
+ * Per-index config for storage-attached index writers.
+ */
+public class IndexWriterConfig
+{
+ public static final String MAXIMUM_NODE_CONNECTIONS = "maximum_node_connections";
Review Comment:
I think the constants related to the same property could be grouped together for easier reading, as in:
```java
public static final String MAXIMUM_NODE_CONNECTIONS = "maximum_node_connections";
public static final int MAXIMUM_MAXIMUM_NODE_CONNECTIONS = 512;
public static final int DEFAULT_MAXIMUM_NODE_CONNECTIONS = 16;
public static final String CONSTRUCTION_BEAM_WIDTH = "construction_beam_width";
public static final int MAXIMUM_CONSTRUCTION_BEAM_WIDTH = 3200;
public static final int DEFAULT_CONSTRUCTION_BEAM_WIDTH = 100;
```
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/VectorIndexSegmentSearcher.java:
##########
@@ -0,0 +1,283 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import javax.annotation.Nullable;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.agrona.collections.IntArrayList;
+import org.apache.cassandra.db.PartitionPosition;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.postings.ReorderingPostingList;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.IntArrayPostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.utils.RangeUtil;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.SparseFixedBitSet;
+
+/**
+ * Executes ann search against the HNSW graph for an individual index segment.
+ */
+public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
+{
+ private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+
+ private final CassandraOnDiskHnsw graph;
+ private final VectorType<float[]> type;
+ private int maxBruteForceRows; // not final so test can inject its own setting
+ private final ThreadLocal<SparseFixedBitSet> cachedBitSets;
+
+ VectorIndexSegmentSearcher(PrimaryKeyMap.Factory primaryKeyMapFactory,
+ PerColumnIndexFiles perIndexFiles,
+ SegmentMetadata segmentMetadata,
+ IndexContext indexContext) throws IOException
+ {
+ super(primaryKeyMapFactory, perIndexFiles, segmentMetadata, indexContext);
+ graph = new CassandraOnDiskHnsw(segmentMetadata.componentMetadatas, perIndexFiles, indexContext);
+ type = (VectorType<float[]>) indexContext.getValidator();
Review Comment:
This needs `@SuppressWarnings("unchecked")`
##########
src/java/org/apache/cassandra/index/Index.java:
##########
@@ -440,6 +445,19 @@ default boolean filtersMultipleContains()
*/
public RowFilter getPostIndexQueryFilter(RowFilter filter);
+ /**
+ * Return a comparator that reorders query result before sending to client
+ *
+ * @param restriction restriction that requires current index
+ * @param columnIndex idx of the indexed column in returned row
+ * @param options query options
+ * @return comparator that for post-query ordering; or null if not supported
Review Comment:
```suggestion
* @return a comparator for post-query ordering; or null if not supported
```
##########
src/java/org/apache/cassandra/index/sai/VectorContext.java:
##########
@@ -0,0 +1,168 @@
+/*
+ * 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.index.sai;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnHeapHnsw;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.lucene.util.Bits;
+
+public class VectorContext
+{
+ public int hnswVectorsAccessed;
+ public int hnswVectorCacheHits;
+
+ private TreeSet<PrimaryKey> shadowedPrimaryKeys; // allocate when needed
+
+ public void recordShadowedPrimaryKey(PrimaryKey primaryKey)
+ {
+ if (shadowedPrimaryKeys == null)
+ shadowedPrimaryKeys = new TreeSet<>();
+ shadowedPrimaryKeys.add(primaryKey);
+ }
+
+ // Returns true if the row ID will be included or false if the row ID will be shadowed
+ public boolean shouldInclude(long sstableRowId, PrimaryKeyMap primaryKeyMap)
+ {
+ return shadowedPrimaryKeys == null || !shadowedPrimaryKeys.contains(primaryKeyMap.primaryKeyFromRowId(sstableRowId));
+ }
+
+ public boolean containsShadowedPrimaryKey(PrimaryKey primaryKey)
+ {
+ return shadowedPrimaryKeys != null && shadowedPrimaryKeys.contains(primaryKey);
+ }
+
+ /**
+ * @return shadowed primary keys, in ascending order
+ */
+ public NavigableSet<PrimaryKey> getShadowedPrimaryKeys()
+ {
+ if (shadowedPrimaryKeys == null)
+ return Collections.emptyNavigableSet();
+ return shadowedPrimaryKeys;
+ }
+
+ public Bits bitsetForShadowedPrimaryKeys(CassandraOnHeapHnsw<PrimaryKey> graph)
+ {
+ if (shadowedPrimaryKeys == null)
+ return null;
+
+ return new IgnoredKeysBits(graph, shadowedPrimaryKeys);
+ }
+
+ public Bits bitsetForShadowedPrimaryKeys(SegmentMetadata metadata, PrimaryKeyMap primaryKeyMap, CassandraOnDiskHnsw graph) throws IOException
+ {
+ Set<Integer> ignoredOrdinals = null;
+ try (var ordinalsView = graph.getOrdinalsView())
+ {
+ for (PrimaryKey primaryKey : getShadowedPrimaryKeys())
+ {
+ // not in current segment
+ if (primaryKey.compareTo(metadata.minKey) < 0 || primaryKey.compareTo(metadata.maxKey) > 0)
+ continue;
+
+ long sstableRowId = primaryKeyMap.rowIdFromPrimaryKey(primaryKey);
+ if (sstableRowId == Long.MAX_VALUE) // not found
+ continue;
+
+ int segmentRowId = Math.toIntExact(sstableRowId - metadata.rowIdOffset);
+ // not in segment yet
+ if (segmentRowId < 0)
+ continue;
+ // end of segment
+ if (segmentRowId > metadata.maxSSTableRowId)
+ break;
+
+ int ordinal = ordinalsView.getOrdinalForRowId(segmentRowId);
+ if (ordinal >= 0)
+ {
+ if (ignoredOrdinals == null)
+ ignoredOrdinals = new HashSet<>();
+ ignoredOrdinals.add(ordinal);
+ }
+ }
+ }
+
+ if (ignoredOrdinals == null)
+ return null;
+
+ return new IgnoringBits(graph, ignoredOrdinals, metadata);
+ }
+
+ private static class IgnoringBits implements Bits
+ {
+ private final CassandraOnDiskHnsw graph;
Review Comment:
It seems we are not using `graph` for anything, can it be removed?
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +73,28 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row Id for a given {@link Token}
+ *
+ * @param range the {@link AbstractBounds} to lookup
+ * @return the first row Id associated with the {@link AbstractBounds}
Review Comment:
```suggestion
* Returns the first row ID for a given partition range.
*
* @param range the {@link AbstractBounds} to lookup
* @return the first row ID associated with the {@link AbstractBounds}
```
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java:
##########
@@ -0,0 +1,48 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.PostingList;
+
+/**
+ * There are two steps in ordering:
+ * <p>
+ * 1. Limit a single sstable's results to the correct keys. At this stage
+ * we put them back in primary key order to play nice with the rest of
+ * the query pipeline.
+ * 2. Merge the results from multiple sstables. Now we leave them in the
+ * final, correct order.
+ * <p>
+ * SegmentOrdering handles the first step.
Review Comment:
Nice JavaDoc. It could be helpful to include a hint on who handles the second step.
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/VectorIndexSegmentSearcher.java:
##########
@@ -0,0 +1,283 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import javax.annotation.Nullable;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.agrona.collections.IntArrayList;
+import org.apache.cassandra.db.PartitionPosition;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.postings.ReorderingPostingList;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.IntArrayPostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.utils.RangeUtil;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.SparseFixedBitSet;
+
+/**
+ * Executes ann search against the HNSW graph for an individual index segment.
+ */
+public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
+{
+ private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+
+ private final CassandraOnDiskHnsw graph;
+ private final VectorType<float[]> type;
+ private int maxBruteForceRows; // not final so test can inject its own setting
Review Comment:
Nit: can be `final`
##########
src/java/org/apache/cassandra/index/sai/disk/v1/postings/ReorderingPostingList.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.index.sai.disk.v1.postings;
+
+import java.io.IOException;
+import java.util.PrimitiveIterator;
+
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.lucene.util.LongHeap;
+
+public class ReorderingPostingList implements PostingList
Review Comment:
Can we add some JavaDoc saying what this posting list implementation does?
##########
src/java/org/apache/cassandra/index/sai/disk/v1/postings/PostingListRangeIterator.java:
##########
@@ -83,10 +82,23 @@ public PostingListRangeIterator(IndexContext indexContext,
this.indexContext = indexContext;
this.primaryKeyMap = primaryKeyMap;
this.postingList = searcherContext.postingList;
- this.searcherContext = searcherContext;
+ this.rowIdOffset = searcherContext.segmentRowIdOffset;
this.queryContext = searcherContext.context;
}
+ public PostingListRangeIterator(IndexContext indexContext,
+ PrimaryKeyMap primaryKeyMap,
+ PostingList postingList,
+ QueryContext queryContext)
+ {
+ super(primaryKeyMap.primaryKeyFromRowId(postingList.minimum()), primaryKeyMap.primaryKeyFromRowId(postingList.maximum()), postingList.size());
Review Comment:
Nit: break the long line
```suggestion
super(primaryKeyMap.primaryKeyFromRowId(postingList.minimum()),
primaryKeyMap.primaryKeyFromRowId(postingList.maximum()),
postingList.size());
```
##########
src/java/org/apache/cassandra/index/sai/disk/v1/IndexWriterConfig.java:
##########
@@ -0,0 +1,164 @@
+/*
+ * 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.index.sai.disk.v1;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.apache.cassandra.config.CassandraRelevantProperties;
+import org.apache.cassandra.db.marshal.AbstractType;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+import static org.apache.cassandra.config.CassandraRelevantProperties.SAI_VECTOR_SEARCH_MAX_TOP_K;
+
+/**
+ * Per-index config for storage-attached index writers.
+ */
+public class IndexWriterConfig
+{
+ public static final String MAXIMUM_NODE_CONNECTIONS = "maximum_node_connections";
+ public static final String CONSTRUCTION_BEAM_WIDTH = "construction_beam_width";
+ public static final String SIMILARITY_FUNCTION = "similarity_function";
+
+ public static final int MAXIMUM_MAXIMUM_NODE_CONNECTIONS = 512;
+ public static final int MAXIMUM_CONSTRUCTION_BEAM_WIDTH = 3200;
+
+ public static final int DEFAULT_MAXIMUM_NODE_CONNECTIONS = 16;
+ public static final int DEFAULT_CONSTRUCTION_BEAM_WIDTH = 100;
+
+ public static final int MAX_TOP_K = SAI_VECTOR_SEARCH_MAX_TOP_K.getInt();
+
+ public static final VectorSimilarityFunction DEFAULT_SIMILARITY_FUNCTION = VectorSimilarityFunction.COSINE;
+
+ public static final String validSimilarityFunctions = Arrays.stream(VectorSimilarityFunction.values())
+ .map(Enum::name)
+ .collect(Collectors.joining(", "));
+
+ private static final IndexWriterConfig EMPTY_CONFIG = new IndexWriterConfig(-1, -1, null);
+
+ private final int maximumNodeConnections;
+
+ private final int constructionBeamWidth;
+
+ private final VectorSimilarityFunction similarityFunction;
Review Comment:
Can we add some JavaDoc for these attributes?
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +73,28 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row Id for a given {@link Token}
+ *
+ * @param range the {@link AbstractBounds} to lookup
+ * @return the first row Id associated with the {@link AbstractBounds}
+ */
+ default long firstRowIdForRange(AbstractBounds<PartitionPosition> range)
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ /**
+ * Returns the last row Id for a given {@link Token}
+ *
+ * @param range the {@link AbstractBounds} to lookup
+ * @return the last row Id associated with the {@link AbstractBounds}
Review Comment:
```suggestion
* Returns the last row ID for a given partition range.
*
* @param range the {@link AbstractBounds} to lookup
* @return the last row ID associated with the {@link AbstractBounds}
```
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/VectorIndexSegmentSearcher.java:
##########
@@ -0,0 +1,283 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import javax.annotation.Nullable;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.agrona.collections.IntArrayList;
+import org.apache.cassandra.db.PartitionPosition;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.postings.ReorderingPostingList;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.IntArrayPostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.utils.RangeUtil;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.SparseFixedBitSet;
+
+/**
+ * Executes ann search against the HNSW graph for an individual index segment.
Review Comment:
```suggestion
* Executes ANN search against the HNSW graph for an individual index segment.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357115095
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/VectorIndexSegmentSearcher.java:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import javax.annotation.Nullable;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.util.SparseFixedBitSet;
+import org.agrona.collections.IntArrayList;
+import org.apache.cassandra.db.PartitionPosition;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.VectorQueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
+import org.apache.cassandra.index.sai.disk.v1.postings.ReorderingPostingList;
+import org.apache.cassandra.index.sai.disk.v1.vector.DiskAnn;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.IntArrayPostingList;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.utils.RangeUtil;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+
+/**
+ * Executes ANN search against a vector graph for an individual index segment.
+ */
+public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
+{
+ private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+
+ private final DiskAnn graph;
+ private final VectorType<float[]> type;
+ private final int maxBruteForceRows;
+ private final ThreadLocal<SparseFixedBitSet> cachedBitSets;
+
+ @SuppressWarnings("unchecked")
+ VectorIndexSegmentSearcher(PrimaryKeyMap.Factory primaryKeyMapFactory,
+ PerColumnIndexFiles perIndexFiles,
+ SegmentMetadata segmentMetadata,
+ IndexContext indexContext) throws IOException
+ {
+ super(primaryKeyMapFactory, perIndexFiles, segmentMetadata, indexContext);
+ graph = new DiskAnn(segmentMetadata.componentMetadatas, perIndexFiles, indexContext);
+ type = (VectorType<float[]>) indexContext.getValidator();
+ cachedBitSets = ThreadLocal.withInitial(() -> new SparseFixedBitSet(graph.size()));
+
+ // estimate the number of comparisons that a search would require; use brute force if we have
+ // fewer rows involved than that
+ maxBruteForceRows = (int)(indexContext.getIndexWriterConfig().getMaximumNodeConnections() * Math.log(graph.size()));
Review Comment:
this turns out to be a very bad estimate, and estimate quality is much more important than I thought at first, because if you only have M results from the other predicates and you are looking for top K in an index of N, where M << N, the dominant factor in your search is just looking for M valid results.
unfortunately I don't have a single-commit fix for this, it evolved across several in the vsearch branch.
okay to leave this for followup but it's quite important to have it in the released version.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357145878
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/OnHeapGraph.java:
##########
@@ -0,0 +1,360 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.cliffc.high_scale_lib.NonBlockingHashMapLong;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.github.jbellis.jvector.disk.CompressedVectors;
+import io.github.jbellis.jvector.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.graph.GraphIndexBuilder;
+import io.github.jbellis.jvector.graph.GraphSearcher;
+import io.github.jbellis.jvector.pq.ProductQuantization;
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.vector.VectorEncoding;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.db.marshal.AbstractType;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.disk.format.IndexComponent;
+import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
+import org.apache.cassandra.index.sai.disk.io.IndexFileUtils;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.io.util.SequentialWriter;
+import org.apache.cassandra.utils.ByteBufferUtil;
+import org.apache.lucene.util.StringHelper;
+
+public class OnHeapGraph<T>
+{
+ private static final Logger logger = LoggerFactory.getLogger(OnHeapGraph.class);
+
+ private final RamAwareVectorValues vectorValues;
+ private final GraphIndexBuilder<float[]> builder;
+ private final VectorType<?> vectorType;
+ private final VectorSimilarityFunction similarityFunction;
+ private final Map<float[], VectorPostings<T>> postingsMap;
+ private final NonBlockingHashMapLong<VectorPostings<T>> postingsByOrdinal;
+ private final AtomicInteger nextOrdinal = new AtomicInteger();
+ private volatile boolean hasDeletions;
+
+ /**
+ * @param termComparator the vector type
+ * @param indexWriterConfig
+ *
+ * Will create a concurrent object.
+ */
+ public OnHeapGraph(AbstractType<?> termComparator, IndexWriterConfig indexWriterConfig)
+ {
+ this(termComparator, indexWriterConfig, true);
+ }
+
+ /**
+ * @param termComparator the vector type
+ * @param indexWriterConfig the {@link IndexWriterConfig} for the graph
+ * @param concurrent should be true for memtables, false for compaction. Concurrent allows us to search
+ * while building the graph; non-concurrent allows us to avoid synchronization costs.
+ */
+ @SuppressWarnings("unchecked")
+ public OnHeapGraph(AbstractType<?> termComparator, IndexWriterConfig indexWriterConfig, boolean concurrent)
+ {
+ this.vectorType = (VectorType<?>) termComparator;
+ vectorValues = concurrent
+ ? new ConcurrentVectorValues(((VectorType<?>) termComparator).dimension)
+ : new CompactionVectorValues(((VectorType<Float>) termComparator));
+ similarityFunction = indexWriterConfig.getSimilarityFunction();
+ // We need to be able to inexpensively distinguish different vectors, with a slower path
+ // that identifies vectors that are equal but not the same reference. A comparison
+ // based Map (which only needs to look at vector elements until a difference is found)
+ // is thus a better option than hash-based (which has to look at all elements to compute the hash).
+ postingsMap = new ConcurrentSkipListMap<>(Arrays::compare);
+ postingsByOrdinal = new NonBlockingHashMapLong<>();
+
+ builder = new GraphIndexBuilder<>(vectorValues,
+ VectorEncoding.FLOAT32,
+ similarityFunction,
+ indexWriterConfig.getMaximumNodeConnections(),
+ indexWriterConfig.getConstructionBeamWidth(),
+ 1.2f,
+ 1.4f);
+ }
+
+ public int size()
+ {
+ return vectorValues.size();
+ }
+
+ public boolean isEmpty()
+ {
+ return postingsMap.values().stream().allMatch(VectorPostings::isEmpty);
+ }
+
+ /**
+ * @return the incremental bytes ysed by adding the given vector to the index
+ */
+ public long add(ByteBuffer term, T key, InvalidVectorBehavior behavior)
+ {
+ assert term != null && term.remaining() != 0;
+
+ var vector = vectorType.composeAsFloat(term);
+ if (behavior == InvalidVectorBehavior.IGNORE)
+ {
+ try
+ {
+ validateIndexable(vector, similarityFunction);
+ }
+ catch (InvalidRequestException e)
+ {
+ logger.trace("Ignoring invalid vector during index build against existing data: {}", vector, e);
+ return 0;
+ }
+ }
+ else
+ {
+ assert behavior == InvalidVectorBehavior.FAIL;
+ validateIndexable(vector, similarityFunction);
+ }
+
+ var bytesUsed = new AtomicLong();
+ var newVector = new AtomicBoolean();
+ // if the vector is already in the graph, all that happens is that the postings list is updated
Review Comment:
this has a pretty nasty race condition when multiple rows are added concurrently with the same vector. fix is in vsearch 9e3ca22c
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356657246
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/CassandraOnHeapGraph.java:
##########
@@ -0,0 +1,360 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.cliffc.high_scale_lib.NonBlockingHashMapLong;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.github.jbellis.jvector.disk.CompressedVectors;
+import io.github.jbellis.jvector.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.graph.GraphIndexBuilder;
+import io.github.jbellis.jvector.graph.GraphSearcher;
+import io.github.jbellis.jvector.pq.ProductQuantization;
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.vector.VectorEncoding;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.db.marshal.AbstractType;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.disk.format.IndexComponent;
+import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
+import org.apache.cassandra.index.sai.disk.io.IndexFileUtils;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.io.util.SequentialWriter;
+import org.apache.cassandra.utils.ByteBufferUtil;
+import org.apache.lucene.util.StringHelper;
+
+public class CassandraOnHeapGraph<T>
+{
+ private static final Logger logger = LoggerFactory.getLogger(CassandraOnHeapGraph.class);
+
+ private final RamAwareVectorValues vectorValues;
+ private final GraphIndexBuilder<float[]> builder;
+ private final VectorType<?> termComparator;
Review Comment:
Nit: maybe this could be named `vectorType`, `termType` or `type`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357108018
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java:
##########
@@ -0,0 +1,56 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+
+/**
+ * There are two steps in ordering:
+ * <p>
+ * 1. Limit a single sstable's results to the correct keys. At this stage
+ * we put them back in primary key order to play nice with the rest of
+ * the query pipeline.
+ * 2. Merge the results from multiple sstables. Now we leave them in the
+ * final, correct order.
+ * <p>
+ * SegmentOrdering handles the first step.
+ *
+ * The second step is handled during post-reconciliation processing when the
+ * {@link org.apache.cassandra.index.Index.QueryPlan#postProcessor(ReadCommand)} is called by
+ * {@link org.apache.cassandra.db.PartitionRangeReadCommand#postReconciliationProcessing(PartitionIterator)}. If the
+ * {@link org.apache.cassandra.index.Index.QueryPlan} is for a topK query then a reordering filter is returned to return
+ * the partitions in the correct order.
+ */
+public interface SegmentOrdering
+{
+ /**
+ * Reorder, limit, and put back into original order the results from a single sstable
Review Comment:
also this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1363743988
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1295,8 +1378,10 @@ private Term prepareLimit(VariableSpecifications boundNames, Term.Raw limit,
return prepLimit;
}
- private static void verifyOrderingIsAllowed(StatementRestrictions restrictions) throws InvalidRequestException
+ private static void verifyOrderingIsAllowed(StatementRestrictions restrictions, Map<ColumnMetadata, Ordering> orderingColumns) throws InvalidRequestException
Review Comment:
When users try to use not-ANN `ORDER BY` they get an error that says `ORDER BY with 2ndary indexes is not supported`. However, now we will support `ORDER BY ... ANN OF` with 2ndary indexes.
So the message should probably say something like `ORDER BY with 2ndary indexes is not supported, except for ANN queries`, or something like that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1373032732
##########
src/java/org/apache/cassandra/service/reads/range/ScanAllRangesCommandIterator.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.service.reads.range;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import com.google.common.base.Preconditions;
+
+import org.apache.cassandra.db.ConsistencyLevel;
+import org.apache.cassandra.db.EmptyIterators;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.PartitionRangeReadCommand;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.locator.EndpointsForRange;
+import org.apache.cassandra.locator.InetAddressAndPort;
+import org.apache.cassandra.locator.ReplicaPlan;
+import org.apache.cassandra.locator.ReplicaPlans;
+import org.apache.cassandra.net.Message;
+import org.apache.cassandra.net.MessagingService;
+import org.apache.cassandra.service.reads.DataResolver;
+import org.apache.cassandra.service.reads.ReadCallback;
+import org.apache.cassandra.service.reads.repair.NoopReadRepair;
+import org.apache.cassandra.tracing.Tracing;
+import org.apache.cassandra.utils.CloseableIterator;
+
+/**
+ * A custom {@link RangeCommandIterator} that queries all replicas required by consistency level at once with data range
+ * specify in {@link PartitionRangeReadCommand}.
+ *
+ * This is to speed up {@link Index.QueryPlan#isTopK()} queries that needs to find global top-k rows in the cluster, because
+ * existing {@link RangeCommandIterator} has to execute a top-k search per vnode range which is wasting resources.
+ */
+public class ScanAllRangesCommandIterator extends RangeCommandIterator
Review Comment:
I think we are missing that comment on `ScanAllRangesCommandIterator` saying that it can only be used with CL=ONE/LOCAL_ONE.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "JeremiahDJordan (via GitHub)" <gi...@apache.org>.
JeremiahDJordan commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370412676
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +71,22 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row ID for a given token or a negative value if not found
+ *
+ * @param token the {@link Token} to lookup
+ * @return the first row ID associated with the {@link Token} or a negative value
+ */
+ long ceiling(Token token);
+
+ /**
+ * Returns the last row ID for a given partition range or the previous row ID if there was no match for the partition range
Review Comment:
My understanding is that ceiling is returning the first row id following a given token, "round up", and floor is returning the first row id before a given token, "round down", which is the last row id of the previous token. I can see how there can be slight confusion on first and last here with floor and ceiling. My understanding is given `a b c X d e f`. `floor(X)` returns `c` and `ceiling(X)` returns `d`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "michaelsembwever (via GitHub)" <gi...@apache.org>.
michaelsembwever commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1375405291
##########
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java:
##########
@@ -455,6 +456,18 @@ public enum CassandraRelevantProperties
SAI_TEST_BALANCED_TREE_DEBUG_ENABLED("cassandra.sai.test.balanced_tree_debug_enabled", "false"),
SAI_TEST_DISABLE_TIMEOUT("cassandra.sai.test.disable.timeout", "false"),
+ /** Whether to allow the user to specify custom options to the hnsw index */
Review Comment:
are we not now DiskANN index ?
(this can be ninja-fixed l8r)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358474705
##########
src/java/org/apache/cassandra/index/sai/disk/format/IndexComponent.java:
##########
@@ -50,6 +50,11 @@ public enum IndexComponent
*/
TERMS_DATA("TermsData"),
+ /**
+ * Product Quantization store used to store compressed vectors for the vector index
+ */
+ VECTORS("Vectors"),
Review Comment:
At one point we had this as `PRODUCT_QUANTIZATION`, but that seemed to be meaningless to anyone. I think `COMPRESSED_VECTORS` is clearest, even though it is a bit long. `PQ_VECTORS` is shorter but isn't clear as to what it contains (unless you know).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357099186
##########
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java:
##########
@@ -424,8 +424,9 @@ public enum CassandraRelevantProperties
SAI_INTERSECTION_CLAUSE_LIMIT("cassandra.sai.intersection_clause_limit", "2"),
/** Latest version to be used for SAI index writing */
SAI_LATEST_VERSION("cassandra.sai.latest_version", "aa"),
- SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size_kb", "5"),
- SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size_kb", "1"),
+ SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size", "5KiB"),
+ SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size", "1KiB"),
+ SAI_MAX_VECTOR_TERM_SIZE("cassandra.sai.max_vector_term_size", "16KiB"),
Review Comment:
imo KB makes more sense so it aligns w/ the other term size limits
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357149666
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/OnHeapGraph.java:
##########
@@ -0,0 +1,360 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.cliffc.high_scale_lib.NonBlockingHashMapLong;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.github.jbellis.jvector.disk.CompressedVectors;
+import io.github.jbellis.jvector.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.graph.GraphIndexBuilder;
+import io.github.jbellis.jvector.graph.GraphSearcher;
+import io.github.jbellis.jvector.pq.ProductQuantization;
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.vector.VectorEncoding;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.db.marshal.AbstractType;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.disk.format.IndexComponent;
+import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
+import org.apache.cassandra.index.sai.disk.io.IndexFileUtils;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.io.util.SequentialWriter;
+import org.apache.cassandra.utils.ByteBufferUtil;
+import org.apache.lucene.util.StringHelper;
+
+public class OnHeapGraph<T>
+{
+ private static final Logger logger = LoggerFactory.getLogger(OnHeapGraph.class);
+
+ private final RamAwareVectorValues vectorValues;
+ private final GraphIndexBuilder<float[]> builder;
+ private final VectorType<?> vectorType;
+ private final VectorSimilarityFunction similarityFunction;
+ private final Map<float[], VectorPostings<T>> postingsMap;
+ private final NonBlockingHashMapLong<VectorPostings<T>> postingsByOrdinal;
+ private final AtomicInteger nextOrdinal = new AtomicInteger();
+ private volatile boolean hasDeletions;
+
+ /**
+ * @param termComparator the vector type
+ * @param indexWriterConfig
+ *
+ * Will create a concurrent object.
+ */
+ public OnHeapGraph(AbstractType<?> termComparator, IndexWriterConfig indexWriterConfig)
+ {
+ this(termComparator, indexWriterConfig, true);
+ }
+
+ /**
+ * @param termComparator the vector type
+ * @param indexWriterConfig the {@link IndexWriterConfig} for the graph
+ * @param concurrent should be true for memtables, false for compaction. Concurrent allows us to search
+ * while building the graph; non-concurrent allows us to avoid synchronization costs.
+ */
+ @SuppressWarnings("unchecked")
+ public OnHeapGraph(AbstractType<?> termComparator, IndexWriterConfig indexWriterConfig, boolean concurrent)
+ {
+ this.vectorType = (VectorType<?>) termComparator;
+ vectorValues = concurrent
+ ? new ConcurrentVectorValues(((VectorType<?>) termComparator).dimension)
+ : new CompactionVectorValues(((VectorType<Float>) termComparator));
+ similarityFunction = indexWriterConfig.getSimilarityFunction();
+ // We need to be able to inexpensively distinguish different vectors, with a slower path
+ // that identifies vectors that are equal but not the same reference. A comparison
+ // based Map (which only needs to look at vector elements until a difference is found)
+ // is thus a better option than hash-based (which has to look at all elements to compute the hash).
+ postingsMap = new ConcurrentSkipListMap<>(Arrays::compare);
+ postingsByOrdinal = new NonBlockingHashMapLong<>();
+
+ builder = new GraphIndexBuilder<>(vectorValues,
+ VectorEncoding.FLOAT32,
+ similarityFunction,
+ indexWriterConfig.getMaximumNodeConnections(),
+ indexWriterConfig.getConstructionBeamWidth(),
+ 1.2f,
+ 1.4f);
+ }
+
+ public int size()
+ {
+ return vectorValues.size();
+ }
+
+ public boolean isEmpty()
+ {
+ return postingsMap.values().stream().allMatch(VectorPostings::isEmpty);
+ }
+
+ /**
+ * @return the incremental bytes ysed by adding the given vector to the index
+ */
+ public long add(ByteBuffer term, T key, InvalidVectorBehavior behavior)
+ {
+ assert term != null && term.remaining() != 0;
+
+ var vector = vectorType.composeAsFloat(term);
+ if (behavior == InvalidVectorBehavior.IGNORE)
+ {
+ try
+ {
+ validateIndexable(vector, similarityFunction);
+ }
+ catch (InvalidRequestException e)
+ {
+ logger.trace("Ignoring invalid vector during index build against existing data: {}", vector, e);
+ return 0;
+ }
+ }
+ else
+ {
+ assert behavior == InvalidVectorBehavior.FAIL;
+ validateIndexable(vector, similarityFunction);
+ }
+
+ var bytesUsed = new AtomicLong();
+ var newVector = new AtomicBoolean();
+ // if the vector is already in the graph, all that happens is that the postings list is updated
+ // otherwise, we add the vector in this order:
+ // 1. to the vectorValues
+ // 2. to the postingsMap
+ // 3. to the graph
+ // This way, concurrent searches of the graph won't see the vector until it's visible
+ // in the other structures as well.
+ var postings = postingsMap.computeIfAbsent(vector, v -> {
+ var bytes = RamEstimation.concurrentHashMapRamUsed(1); // the new posting Map entry
+ var ordinal = nextOrdinal.getAndIncrement();
+ bytes += (vectorValues instanceof ConcurrentVectorValues)
+ ? ((ConcurrentVectorValues) vectorValues).add(ordinal, vector)
+ : ((CompactionVectorValues) vectorValues).add(ordinal, term);
+ bytes += VectorPostings.emptyBytesUsed();
+ bytesUsed.addAndGet(bytes);
+ newVector.set(true);
+ var vp = new VectorPostings<T>(ordinal);
+ postingsByOrdinal.put(ordinal, vp);
+ return vp;
+ });
+ if (postings.add(key))
+ {
+ bytesUsed.addAndGet(VectorPostings.bytesPerPosting());
+ if (newVector.get()) {
+ bytesUsed.addAndGet(builder.addGraphNode(postings.getOrdinal(), vectorValues));
+ }
+ }
+ return bytesUsed.get();
+ }
+
+ // copied out of a Lucene PR -- hopefully committed soon
+ public static final float MAX_FLOAT32_COMPONENT = 1E17f;
+
+ public static void checkInBounds(float[] v)
+ {
+ for (int i = 0; i < v.length; i++)
+ {
+ if (!Float.isFinite(v[i]))
+ {
+ throw new IllegalArgumentException("non-finite value at vector[" + i + "]=" + v[i]);
+ }
+
+ if (Math.abs(v[i]) > MAX_FLOAT32_COMPONENT)
+ {
+ throw new IllegalArgumentException("Out-of-bounds value at vector[" + i + "]=" + v[i]);
+ }
+ }
+ }
+
+ public static void validateIndexable(float[] vector, VectorSimilarityFunction similarityFunction)
+ {
+ try
+ {
+ checkInBounds(vector);
+ }
+ catch (IllegalArgumentException e)
+ {
+ throw new InvalidRequestException(e.getMessage());
+ }
+
+ if (similarityFunction == VectorSimilarityFunction.COSINE)
+ {
+ for (int i = 0; i < vector.length; i++)
+ {
+ if (vector[i] != 0)
+ return;
+ }
+ throw new InvalidRequestException("Zero vectors cannot be indexed or queried with cosine similarity");
+ }
+ }
+
+ public Collection<T> keysFromOrdinal(int node)
+ {
+ return postingsByOrdinal.get(node).getPostings();
+ }
+
+ public long remove(ByteBuffer term, T key)
+ {
+ assert term != null && term.remaining() != 0;
+
+ var vector = vectorType.composeAsFloat(term);
+ var postings = postingsMap.get(vector);
+ if (postings == null)
+ {
+ // it's possible for this to be called against a different memtable than the one
+ // the value was originally added to, in which case we do not expect to find
+ // the key among the postings for this vector
+ return 0;
+ }
+
+ hasDeletions = true;
+ return postings.remove(key);
+ }
+
+ /**
+ * @return keys (PrimaryKey or segment row id) associated with the topK vectors near the query
+ */
+ public PriorityQueue<T> search(float[] queryVector, int limit, Bits toAccept)
+ {
+ validateIndexable(queryVector, similarityFunction);
+
+ // search() errors out when an empty graph is passed to it
+ if (vectorValues.size() == 0)
+ return new PriorityQueue<>();
+
+ Bits bits = hasDeletions ? BitsUtil.bitsIgnoringDeleted(toAccept, postingsByOrdinal) : toAccept;
+ // VSTODO re-use searcher objects
+ var result = GraphSearcher.search(queryVector,
+ limit,
+ vectorValues,
+ VectorEncoding.FLOAT32,
+ similarityFunction,
+ builder.getGraph(),
+ bits);
+ var a = result.getNodes();
+ PriorityQueue<T> keyQueue = new PriorityQueue<>();
+ for (int i = 0; i < a.length; i++)
+ keyQueue.addAll(keysFromOrdinal(a[i].node));
+ return keyQueue;
+ }
+
+ public SegmentMetadata.ComponentMetadataMap writeData(IndexDescriptor indexDescriptor, IndexContext indexContext, Function<T, Integer> postingTransformer) throws IOException
+ {
+ int nInProgress = builder.insertsInProgress();
+ assert nInProgress == 0 : String.format("Attempting to write graph while %d inserts are in progress", nInProgress);
+ assert nextOrdinal.get() == builder.getGraph().size() : String.format("nextOrdinal %d != graph size %d -- ordinals should be sequential",
+ nextOrdinal.get(), builder.getGraph().size());
+ assert vectorValues.size() == builder.getGraph().size() : String.format("vector count %d != graph size %d",
+ vectorValues.size(), builder.getGraph().size());
+ assert postingsMap.keySet().size() == vectorValues.size() : String.format("postings map entry count %d != vector count %d",
+ postingsMap.keySet().size(), vectorValues.size());
+ logger.debug("Writing graph with {} rows and {} distinct vectors", postingsMap.values().stream().mapToInt(VectorPostings::size).sum(), vectorValues.size());
+
+ try (var pqOutput = IndexFileUtils.instance.openOutput(indexDescriptor.fileFor(IndexComponent.VECTORS, indexContext), true);
+ var postingsOutput = IndexFileUtils.instance.openOutput(indexDescriptor.fileFor(IndexComponent.POSTING_LISTS, indexContext), true);
+ var indexOutput = IndexFileUtils.instance.openOutput(indexDescriptor.fileFor(IndexComponent.TERMS_DATA, indexContext), true))
+ {
+ // compute and write PQ
+ long pqOffset = pqOutput.getFilePointer();
+ long pqPosition = writePQ(pqOutput.asSequentialWriter());
+ long pqLength = pqPosition - pqOffset;
+
+ var deletedOrdinals = new HashSet<Integer>();
+ postingsMap.values().stream().filter(VectorPostings::isEmpty).forEach(vectorPostings -> deletedOrdinals.add(vectorPostings.getOrdinal()));
+ // remove ordinals that don't have corresponding row ids due to partition/range deletion
+ for (VectorPostings<T> vectorPostings : postingsMap.values())
+ {
+ vectorPostings.computeRowIds(postingTransformer);
+ if (vectorPostings.shouldAppendDeletedOrdinal())
+ deletedOrdinals.add(vectorPostings.getOrdinal());
+ }
+ // write postings
+ long postingsOffset = postingsOutput.getFilePointer();
+ long postingsPosition = new VectorPostingsWriter<T>().writePostings(postingsOutput.asSequentialWriter(), vectorValues, postingsMap, deletedOrdinals);
+ long postingsLength = postingsPosition - postingsOffset;
+
+ // complete (internal clean up) and write the graph
+ builder.complete();
+ long termsOffset = indexOutput.getFilePointer();
+ OnDiskGraphIndex.write(builder.getGraph(), vectorValues, indexOutput.asSequentialWriter());
+ long termsLength = indexOutput.getFilePointer() - termsOffset;
+
+ // add components to the metadata map
+ SegmentMetadata.ComponentMetadataMap metadataMap = new SegmentMetadata.ComponentMetadataMap();
+ metadataMap.put(IndexComponent.TERMS_DATA, -1, termsOffset, termsLength, Map.of());
+ metadataMap.put(IndexComponent.POSTING_LISTS, -1, postingsOffset, postingsLength, Map.of());
+ Map<String, String> vectorConfigs = Map.of("SEGMENT_ID", ByteBufferUtil.bytesToHex(ByteBuffer.wrap(StringHelper.randomId())));
+ metadataMap.put(IndexComponent.VECTORS, -1, pqOffset, pqLength, vectorConfigs);
+ return metadataMap;
+ }
+ }
+
+ @SuppressWarnings("SynchronizationOnLocalVariableOrMethodParameter")
+ private long writePQ(SequentialWriter writer) throws IOException
+ {
+ // don't bother with PQ if there are fewer than 1K vectors
+ int M = vectorValues.dimension() / 2;
+ writer.writeBoolean(vectorValues.size() >= 1024);
+ if (vectorValues.size() < 1024)
+ {
+ logger.debug("Skipping PQ for only {} vectors", vectorValues.size());
+ return writer.position();
+ }
+
+ logger.debug("Computing PQ for {} vectors", vectorValues.size());
+ // hack to only do this one at a time during compaction since we're reading a ton of vectors into memory
+ var synchronizeOn = vectorValues instanceof ConcurrentVectorValues ? this : OnHeapGraph.class;
+ synchronized (synchronizeOn)
+ {
+ // train PQ and encode
+ var pq = ProductQuantization.compute(vectorValues, M, false);
+ // VSTODO assert !vectorValues.isValueShared();
+ var encoded = IntStream.range(0, vectorValues.size()).parallel()
+ .mapToObj(i -> pq.encode(vectorValues.vectorValue(i)))
+ .collect(Collectors.toList());
+ var cv = new CompressedVectors(pq, encoded);
Review Comment:
this line and the following ones can move out of the synchronized block
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358480588
##########
src/java/org/apache/cassandra/index/sai/plan/vector/VectorQueryController.java:
##########
@@ -0,0 +1,247 @@
+/*
+ * 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.index.sai.plan.vector;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.memtable.Memtable;
+import org.apache.cassandra.exceptions.QueryCancelledException;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.disk.SSTableRowIdKeyRangeIterator;
+import org.apache.cassandra.index.sai.disk.v1.postings.MergePostingList;
+import org.apache.cassandra.index.sai.disk.v1.vector.CheckpointingIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator;
+import org.apache.cassandra.index.sai.metrics.TableQueryMetrics;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.plan.StorageAttachedIndexQueryController;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.postings.PostingListIntersection;
+import org.apache.cassandra.index.sai.postings.PostingListUnion;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.utils.Throwables;
+
+public class VectorQueryController extends StorageAttachedIndexQueryController
Review Comment:
This will disappear with the rework of the query path. I'll get those changes in asap.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1363983036
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
Review Comment:
We shouldn't mention CC here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1332990976
##########
src/java/org/apache/cassandra/index/Index.java:
##########
@@ -440,6 +445,19 @@ default boolean filtersMultipleContains()
*/
public RowFilter getPostIndexQueryFilter(RowFilter filter);
+ /**
+ * Return a comparator that reorders query result before sending to client
+ *
+ * @param restriction restriction that requires current index
+ * @param columnIndex idx of the indexed column in returned row
+ * @param options query options
+ * @return comparator that for post-query ordering; or null if not supported
+ */
+ default Comparator<List<ByteBuffer>> getPostQueryOrdering(Restriction restriction, int columnIndex, QueryOptions options)
Review Comment:
This assumes that the index implementation orders results by a single columns. I think the index API will be more generic if it can take multiple columns:
```suggestion
default Comparator<List<ByteBuffer>> getPostQueryOrdering(Restriction restriction, QueryOptions options, int... columnIndexes)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341287873
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -40,16 +41,21 @@
import java.util.stream.Collectors;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.Futures; // checkstyle: permit this import
import com.google.common.util.concurrent.ListenableFuture; // checkstyle: permit this import
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
Review Comment:
Thanks, I thought I had caught all of those. I will do a final check before putting this up for review again that I have caught all of these.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341315868
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -284,11 +288,13 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
nowInSec, userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+ checkFalse(query.isTopK() && options.getConsistency() != ConsistencyLevel.ONE, TOPK_CONSISTENCY_LEVEL_ERROR);
Review Comment:
I think we can also allow `LOCAL_ONE`. IIRC that consistency level also uses a single replica, so there shouldn't be reconciliation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341367309
##########
src/java/org/apache/cassandra/index/sai/VectorContext.java:
##########
@@ -0,0 +1,168 @@
+/*
+ * 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.index.sai;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnHeapHnsw;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.lucene.util.Bits;
+
+public class VectorContext
Review Comment:
Yes. I have also renamed this to `VectorQueryContext` to make it clearer as to it's purpose.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1331686334
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -483,7 +486,7 @@ else if (restrictions.keyIsInRelation())
// Please note that the isExhausted state of the pager only gets updated when we've closed the page, so this
// shouldn't be moved inside the 'try' above.
- if (!pager.isExhausted())
+ if (!pager.isExhausted() && !pager.pager.isTopK())
Review Comment:
While testing how paging works, I have randomly found that a range filter on the partition key can throw an exception:
```java
@Test
public void testKeyRestrictions()
{
createTable("CREATE TABLE %s (k int PRIMARY KEY, v vector<float, 1>)");
createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
execute("INSERT INTO %s (k, v) VALUES (1, [1])");
assertRows(execute("SELECT k, v FROM %s WHERE k > 0 LIMIT 4 ALLOW FILTERING"), row(1, vector(1f)));
assertRows(execute("SELECT k, v FROM %s WHERE k = 1 ORDER BY v ANN OF [0] LIMIT 4 ALLOW FILTERING"), row(1, vector(1f)));
assertRows(execute("SELECT k, v FROM %s WHERE k > 0 ORDER BY v ANN OF [0] LIMIT 4 ALLOW FILTERING"), row(1, vector(1f))); // exception!!
}
```
The error for in-memory is:
```
java.lang.AssertionError: Attempt to use memtable index manager on non-indexed context
at org.apache.cassandra.index.sai.IndexContext.getMemtableIndexManager(IndexContext.java:231)
at org.apache.cassandra.index.sai.plan.vector.VectorQueryController.lambda$getIndexQueryResults$1(VectorQueryController.java:89)
at java.base/java.util.stream.ReferencePipeline$7$1.accept(ReferencePipeline.java:271)
at java.base/java.util.ArrayList$ArrayListSpliterator.forEachRemaining(ArrayList.java:1655)
at java.base/java.util.stream.AbstractPipeline.copyInto(AbstractPipeline.java:484)
at java.base/java.util.stream.AbstractPipeline.wrapAndCopyInto(AbstractPipeline.java:474)
at java.base/java.util.stream.ReduceOps$ReduceOp.evaluateSequential(ReduceOps.java:913)
at java.base/java.util.stream.AbstractPipeline.evaluate(AbstractPipeline.java:234)
at java.base/java.util.stream.ReferencePipeline.collect(ReferencePipeline.java:578)
at org.apache.cassandra.index.sai.plan.vector.VectorQueryController.getIndexQueryResults(VectorQueryController.java:90)
at org.apache.cassandra.index.sai.plan.Operation$AndNode.rangeIterator(Operation.java:305)
at org.apache.cassandra.index.sai.plan.Operation.buildIterator(Operation.java:178)
at org.apache.cassandra.index.sai.plan.StorageAttachedIndexSearcher$ResultRetriever.<init>(StorageAttachedIndexSearcher.java:152)
at org.apache.cassandra.index.sai.plan.StorageAttachedIndexSearcher.lambda$search$0(StorageAttachedIndexSearcher.java:111)
at org.apache.cassandra.index.sai.plan.StorageAttachedIndexSearcher.search(StorageAttachedIndexSearcher.java:118)
at org.apache.cassandra.db.ReadCommand.executeLocally(ReadCommand.java:431)
at org.apache.cassandra.db.AbstractReadQuery.executeInternal(AbstractReadQuery.java:65)
at org.apache.cassandra.db.ReadCommand.executeInternal(ReadCommand.java:88)
at org.apache.cassandra.cql3.statements.SelectStatement.executeInternal(SelectStatement.java:543)
at org.apache.cassandra.cql3.statements.SelectStatement.executeLocally(SelectStatement.java:515)
at org.apache.cassandra.cql3.statements.SelectStatement.executeLocally(SelectStatement.java:108)
at org.apache.cassandra.cql3.QueryProcessor.executeInternal(QueryProcessor.java:445)
at org.apache.cassandra.cql3.CQLTester.executeFormattedQuery(CQLTester.java:1595)
at org.apache.cassandra.cql3.CQLTester.execute(CQLTester.java:1574)
at org.apache.cassandra.cql3.validation.operations.CQLVectorTest.testKeyRestrictions(CQLVectorTest.java:69)
at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
at java.base/java.lang.reflect.Method.invoke(Method.java:566)
at org.junit.runners.model.FrameworkMethod$1.runReflectiveCall(FrameworkMethod.java:50)
at org.junit.internal.runners.model.ReflectiveCallable.run(ReflectiveCallable.java:12)
at org.junit.runners.model.FrameworkMethod.invokeExplosively(FrameworkMethod.java:47)
at org.junit.internal.runners.statements.InvokeMethod.evaluate(InvokeMethod.java:17)
at org.junit.internal.runners.statements.RunBefores.evaluate(RunBefores.java:26)
at org.junit.internal.runners.statements.RunAfters.evaluate(RunAfters.java:27)
at org.junit.runners.ParentRunner.runLeaf(ParentRunner.java:325)
at org.junit.runners.BlockJUnit4ClassRunner.runChild(BlockJUnit4ClassRunner.java:78)
at org.junit.runners.BlockJUnit4ClassRunner.runChild(BlockJUnit4ClassRunner.java:57)
at org.junit.runners.ParentRunner$3.run(ParentRunner.java:290)
at org.junit.runners.ParentRunner$1.schedule(ParentRunner.java:71)
at org.junit.runners.ParentRunner.runChildren(ParentRunner.java:288)
at org.junit.runners.ParentRunner.access$000(ParentRunner.java:58)
at org.junit.runners.ParentRunner$2.evaluate(ParentRunner.java:268)
at org.junit.internal.runners.statements.RunBefores.evaluate(RunBefores.java:26)
at org.junit.internal.runners.statements.RunAfters.evaluate(RunAfters.java:27)
at org.junit.runners.ParentRunner.run(ParentRunner.java:363)
at org.junit.runner.JUnitCore.run(JUnitCore.java:137)
at com.intellij.junit4.JUnit4IdeaTestRunner.startRunnerWithArgs(JUnit4IdeaTestRunner.java:69)
at com.intellij.rt.junit.IdeaTestRunner$Repeater$1.execute(IdeaTestRunner.java:38)
at com.intellij.rt.execution.junit.TestsRepeater.repeat(TestsRepeater.java:11)
at com.intellij.rt.junit.IdeaTestRunner$Repeater.startRunnerWithArgs(IdeaTestRunner.java:35)
at com.intellij.rt.junit.JUnitStarter.prepareStreamsAndStart(JUnitStarter.java:232)
at com.intellij.rt.junit.JUnitStarter.main(JUnitStarter.java:55)
```
Whereas if we flush the table it looks like this:
```
java.lang.UnsupportedOperationException
at org.apache.cassandra.io.filesystem.ListenableFileSystem$ListenablePath.toFile(ListenableFileSystem.java:661)
at org.apache.cassandra.io.util.File.toJavaIOFile(File.java:724)
at org.apache.cassandra.index.sai.disk.format.IndexDescriptor.createComponentOnDisk(IndexDescriptor.java:185)
at org.apache.cassandra.index.sai.disk.v1.SSTableComponentsWriter.complete(SSTableComponentsWriter.java:106)
at org.apache.cassandra.index.sai.disk.StorageAttachedIndexWriter.complete(StorageAttachedIndexWriter.java:171)
at com.google.common.collect.ImmutableList.forEach(ImmutableList.java:422)
at org.apache.cassandra.io.sstable.format.SSTableWriter.commit(SSTableWriter.java:235)
at org.apache.cassandra.io.sstable.SimpleSSTableMultiWriter.commit(SimpleSSTableMultiWriter.java:90)
at org.apache.cassandra.db.ColumnFamilyStore$Flush.flushMemtable(ColumnFamilyStore.java:1321)
at org.apache.cassandra.db.ColumnFamilyStore$Flush.run(ColumnFamilyStore.java:1219)
at org.apache.cassandra.concurrent.ExecutionFailure$1.run(ExecutionFailure.java:133)
at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1128)
at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:628)
at io.netty.util.concurrent.FastThreadLocalRunnable.run(FastThreadLocalRunnable.java:30)
at java.base/java.lang.Thread.run(Thread.java:834)
```
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -288,7 +291,7 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
query.trackWarnings();
ResultMessage.Rows rows;
- if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize)))
+ if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize) || query.isTopK()))
Review Comment:
Also, maybe we should throw a client warning if the query's page size is lower than the limit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1338212436
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -535,7 +538,7 @@ public ResultMessage.Rows executeInternal(QueryState state,
try (ReadExecutionController executionController = query.executionController())
{
- if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize)))
+ if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize) || query.isTopK()))
Review Comment:
Yes we should. I have added validation and tests to do this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#issuecomment-1743025475
@adelapena I have addressed (I believe) the remaining review comments and rebased the branch against the latest cassandra-5.0
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1348558477
##########
src/java/org/apache/cassandra/service/reads/range/ScanAllRangesCommandIterator.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.service.reads.range;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import com.google.common.base.Preconditions;
+
+import org.apache.cassandra.db.ConsistencyLevel;
+import org.apache.cassandra.db.EmptyIterators;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.PartitionRangeReadCommand;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.locator.EndpointsForRange;
+import org.apache.cassandra.locator.InetAddressAndPort;
+import org.apache.cassandra.locator.ReplicaPlan;
+import org.apache.cassandra.locator.ReplicaPlans;
+import org.apache.cassandra.net.Message;
+import org.apache.cassandra.net.MessagingService;
+import org.apache.cassandra.service.reads.DataResolver;
+import org.apache.cassandra.service.reads.ReadCallback;
+import org.apache.cassandra.service.reads.repair.NoopReadRepair;
+import org.apache.cassandra.tracing.Tracing;
+import org.apache.cassandra.utils.CloseableIterator;
+
+/**
+ * A custom {@link RangeCommandIterator} that queries all replicas required by consistency level at once with data range
+ * specify in {@link PartitionRangeReadCommand}.
+ *
+ * This is to speed up {@link Index.QueryPlan#isTopK()} queries that needs to find global top-k rows in the cluster, because
+ * existing {@link RangeCommandIterator} has to execute a top-k search per vnode range which is wasting resources.
+ */
+public class ScanAllRangesCommandIterator extends RangeCommandIterator
Review Comment:
I think I would extend the class JavaDoc a bit to mention that the replica plans for each data range are combined into a single shared replica plan, and perhaps mention the consequences for reconciliation. Particularly, that queries with CL=ONE/LOCAL_ONE will use reconciliation.
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -274,10 +292,51 @@ else if (relation.isLIKE())
throw invalidRequest("Non PRIMARY KEY columns found in where clause: %s ",
Joiner.on(", ").join(nonPrimaryKeyColumns));
}
+
+ var annRestriction = Streams.stream(nonPrimaryKeyRestrictions).filter(SingleRestriction::isANN).findFirst();
+ if (annRestriction.isPresent())
+ {
+ // If there is an ANN restriction then it must be for a vector<float, n> column, and it must have an index
+ var annColumn = annRestriction.get().getFirstColumn();
+
+ if (!annColumn.type.isVector() || !(((VectorType<?>)annColumn.type).elementType instanceof FloatType))
+ throw invalidRequest(StatementRestrictions.ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE);
+ if (indexRegistry == null || indexRegistry.listIndexes().stream().noneMatch(i -> i.dependsOn(annColumn)))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE);
+ // We do not allow ANN query filtering using non-indexed columns
+ var nonAnnColumns = Streams.stream(nonPrimaryKeyRestrictions).filter(r -> !r.isANN()).map(r -> r.getFirstColumn()).collect(Collectors.toList());
+ var clusteringColumns = clusteringColumnsRestrictions.getColumnDefinitions();
+ if (!nonAnnColumns.isEmpty() || !clusteringColumns.isEmpty())
+ {
+ var nonIndexedColumns = Stream.concat(nonAnnColumns.stream(), clusteringColumns.stream())
Review Comment:
I think we also need to consider `partitionKeyRestrictions` to exclude restrictions on columns that are part of a multi-column partition key. Those also require `ALLOW FILTERING` and post-filtering.
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorSegmentationTest.java:
##########
@@ -0,0 +1,142 @@
+/*
+ * 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.index.sai.cql;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.Test;
+
+import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.db.marshal.FloatType;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentBuilder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class VectorSegmentationTest extends VectorTester
+{
+ private static final int dimension = 100;
+
+ @Test
+ public void testMultipleSegmentsForCreatingIndex() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+
+ int vectorCount = 100;
+ List<float[]> vectors = new ArrayList<>();
+ for (int row = 0; row < vectorCount; row++)
+ {
+ float[] vector = nextVector();
+ vectors.add(vector);
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", row, vector(vector));
+ }
+
+ flush();
+
+ SegmentBuilder.updateLastValidSegmentRowId(17); // 17 rows per segment
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ int limit = 35;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ @Test
+ public void testMultipleSegmentsForCompaction() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ List<float[]> vectors = new ArrayList<>();
+ int rowsPerSSTable = 10;
+ int sstables = 5;
+ int pk = 0;
+ for (int i = 0; i < sstables; i++)
+ {
+ for (int row = 0; row < rowsPerSSTable; row++)
+ {
+ float[] vector = nextVector();
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", pk++, vector(vector));
+ vectors.add(vector);
+ }
+
+ flush();
+ }
+
+ int limit = 30;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+
+
+ SegmentBuilder.updateLastValidSegmentRowId(11); // 11 rows per segment
+ compact();
+
+ queryVector = nextVector();
+ resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ resultVectors = getVectorsFromResult(resultSet);
+ recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ protected Vector<Float> vector(float[] values)
+ {
+ Float[] floats = new Float[values.length];
+ for (int i = 0; i < values.length; i++)
+ floats[i] = values[i];
+
+ return new Vector<>(floats);
+ }
Review Comment:
I think this method is overriding `CQLTester#vector(float[])` with a mostly identical implementation.
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -154,15 +155,28 @@ private ResolveContext(E replicas)
true,
command.selectsFullPartition(),
enforceStrictLiveness);
+
+ // In case of top-k query, do not trim reconciled rows here because QueryPlan#postProcessor()
+ // needs to compare all rows
+ if (command.isTopK())
+ this.mergedResultCounter.onlyCount();
}
private boolean needsReadRepair()
{
+ // each replica may return different estimated top-K rows, it doesn't mean data is not replicated.
+ if (command.isTopK())
+ return false;
Review Comment:
We don't need this change, since the check right below makes sure RR is off when the number of replicas is greater than one. Indeed, the distributed tests pass without it.
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,28 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
Review Comment:
We probably should add a line about the inability to combine ANN with filtering.
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -154,15 +155,28 @@ private ResolveContext(E replicas)
true,
command.selectsFullPartition(),
enforceStrictLiveness);
+
+ // In case of top-k query, do not trim reconciled rows here because QueryPlan#postProcessor()
+ // needs to compare all rows
+ if (command.isTopK())
+ this.mergedResultCounter.onlyCount();
}
private boolean needsReadRepair()
{
+ // each replica may return different estimated top-K rows, it doesn't mean data is not replicated.
+ if (command.isTopK())
+ return false;
Review Comment:
Ah, I see what's going on here. The methods to skip SRP and RR are based on the number of replicas on the replica plan, which is identical to the number of responses, etc. Queries with CL=ONE/LOCAL_ONE should use a single replica, so the code skips SRP and RR if there is a single replica and thereof there is nothing to reconcile.
However, this patch adds a new `ScanAllRangesCommandIterator` that combines the separate replica plans of each data range into a single shared replica plan. That's why the `DataResolver` can see multiple replicas in top-k queries with CL=ONE/LOCAL_ONE.
I guess we can keep the `command.isTopK` checks, but add some comments about how the replica plan in top-k queries can have multiple replicas independently of the CL.
A thing that can be revisited in the future is how the new `ScanAllRangesCommandIterator` forces the reconciliation of disjoint ranges, which seems wasteful.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357119265
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/VectorIndexSegmentSearcher.java:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import javax.annotation.Nullable;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.util.SparseFixedBitSet;
+import org.agrona.collections.IntArrayList;
+import org.apache.cassandra.db.PartitionPosition;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.VectorQueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
+import org.apache.cassandra.index.sai.disk.v1.postings.ReorderingPostingList;
+import org.apache.cassandra.index.sai.disk.v1.vector.DiskAnn;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.IntArrayPostingList;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.utils.RangeUtil;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+
+/**
+ * Executes ANN search against a vector graph for an individual index segment.
+ */
+public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
+{
+ private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+
+ private final DiskAnn graph;
+ private final VectorType<float[]> type;
+ private final int maxBruteForceRows;
+ private final ThreadLocal<SparseFixedBitSet> cachedBitSets;
+
+ @SuppressWarnings("unchecked")
+ VectorIndexSegmentSearcher(PrimaryKeyMap.Factory primaryKeyMapFactory,
+ PerColumnIndexFiles perIndexFiles,
+ SegmentMetadata segmentMetadata,
+ IndexContext indexContext) throws IOException
+ {
+ super(primaryKeyMapFactory, perIndexFiles, segmentMetadata, indexContext);
+ graph = new DiskAnn(segmentMetadata.componentMetadatas, perIndexFiles, indexContext);
+ type = (VectorType<float[]>) indexContext.getValidator();
+ cachedBitSets = ThreadLocal.withInitial(() -> new SparseFixedBitSet(graph.size()));
+
+ // estimate the number of comparisons that a search would require; use brute force if we have
+ // fewer rows involved than that
+ maxBruteForceRows = (int)(indexContext.getIndexWriterConfig().getMaximumNodeConnections() * Math.log(graph.size()));
+ }
+
+ @Override
+ public long indexFileCacheSize()
+ {
+ return graph.ramBytesUsed();
+ }
+
+ @Override
+ public PostingList search(Expression exp, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
+ {
+ return toRangePostingList(searchPosting(context.vectorContext(), exp, keyRange), context);
+ }
+
+ private PostingList searchPosting(VectorQueryContext context, Expression exp, AbstractBounds<PartitionPosition> keyRange) throws IOException
+ {
+ if (logger.isTraceEnabled())
+ logger.trace(indexContext.logMessage("Searching on expression '{}'..."), exp);
+
+ if (exp.getOp() != Expression.IndexOperator.ANN)
+ throw new IllegalArgumentException(indexContext.logMessage("Unsupported expression during ANN index query: " + exp));
+
+ BitsOrPostingList bitsOrPostingList = bitsOrPostingListForKeyRange(context, keyRange);
+ if (bitsOrPostingList.skipANN())
+ return bitsOrPostingList.postingList();
+
+ ByteBuffer buffer = exp.lower.value.raw;
+ float[] queryVector = TypeUtil.decomposeVector(indexContext, buffer.duplicate());
+ return graph.search(queryVector, context.limit(), bitsOrPostingList.getBits());
+ }
+
+ /**
+ * Return bit set we need to search the graph; otherwise return posting list to bypass the graph
+ */
+ private BitsOrPostingList bitsOrPostingListForKeyRange(VectorQueryContext context, AbstractBounds<PartitionPosition> keyRange) throws IOException
+ {
+ // create a bitset of ordinals corresponding to the rows in the given key range
+ SparseFixedBitSet bits = bitSetForSearch();
+ boolean hasMatches = false;
+
+ try (PrimaryKeyMap primaryKeyMap = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap())
+ {
+ // not restricted
+ if (RangeUtil.coversFullRing(keyRange))
+ return new BitsOrPostingList(context.bitsetForShadowedPrimaryKeys(metadata, primaryKeyMap, graph));
+
+ // it will return the next row id if given key is not found.
+ long minSSTableRowId = primaryKeyMap.firstRowIdForRange(keyRange);
+ long maxSSTableRowId = primaryKeyMap.lastRowIdForRange(keyRange);
+
+ if (minSSTableRowId > maxSSTableRowId)
+ return new BitsOrPostingList(PostingList.EMPTY);
+
+ // if it covers entire segment, skip bit set
+ if (minSSTableRowId <= metadata.minSSTableRowId && maxSSTableRowId >= metadata.maxSSTableRowId)
+ return new BitsOrPostingList(context.bitsetForShadowedPrimaryKeys(metadata, primaryKeyMap, graph));
+
+ minSSTableRowId = Math.max(minSSTableRowId, metadata.minSSTableRowId);
+ maxSSTableRowId = Math.min(maxSSTableRowId, metadata.maxSSTableRowId);
+
+ // if num of matches are not bigger than limit, skip ANN
+ var nRows = maxSSTableRowId - minSSTableRowId + 1;
Review Comment:
this is a bad estimate for nRows. also fixed in vsearch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357140298
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/DiskAnn.java:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.PrimitiveIterator;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import io.github.jbellis.jvector.disk.CachingGraphIndex;
+import io.github.jbellis.jvector.disk.CompressedVectors;
+import io.github.jbellis.jvector.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.disk.ReaderSupplier;
+import io.github.jbellis.jvector.graph.GraphIndex;
+import io.github.jbellis.jvector.graph.GraphSearcher;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.NeighborSimilarity;
+import io.github.jbellis.jvector.graph.SearchResult.NodeScore;
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.disk.format.IndexComponent;
+import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
+import org.apache.cassandra.index.sai.disk.v1.postings.ReorderingPostingList;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.io.util.FileHandle;
+import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.io.util.RandomAccessReader;
+
+public class DiskAnn implements AutoCloseable
+{
+ private final OnDiskOrdinalsMap ordinalsMap;
+ private final CachingGraphIndex graph;
+ private final VectorSimilarityFunction similarityFunction;
+
+ // only one of these will be not null
+ private final CompressedVectors compressedVectors;
+ private final ListRandomAccessVectorValues uncompressedVectors;
+
+ public DiskAnn(SegmentMetadata.ComponentMetadataMap componentMetadatas, PerColumnIndexFiles indexFiles, IndexContext context) throws IOException
+ {
+ similarityFunction = context.getIndexWriterConfig().getSimilarityFunction();
+
+ SegmentMetadata.ComponentMetadata termsMetadata = componentMetadatas.get(IndexComponent.TERMS_DATA);
+ graph = new CachingGraphIndex(new OnDiskGraphIndex<>(new TermsReaderSupplier(indexFiles.termsData()), termsMetadata.offset));
+
+ long pqSegmentOffset = componentMetadatas.get(IndexComponent.VECTORS).offset;
+ try (var fileHandle = indexFiles.vectors(); var reader = new RandomAccessReaderAdapter(fileHandle.createReader()))
+ {
+ reader.seek(pqSegmentOffset);
+ var containsCompressedVectors = reader.readBoolean();
+ if (containsCompressedVectors) {
+ compressedVectors = CompressedVectors.load(reader, reader.getFilePointer());
+ uncompressedVectors = null;
+ }
+ else
+ {
+ compressedVectors = null;
+ var vectors = cacheOriginalVectors(graph);
+ uncompressedVectors = new ListRandomAccessVectorValues(vectors, vectors.get(0).length);
+ }
+ }
+
+ SegmentMetadata.ComponentMetadata postingListsMetadata = componentMetadatas.get(IndexComponent.POSTING_LISTS);
+ ordinalsMap = new OnDiskOrdinalsMap(indexFiles.postingLists(), postingListsMetadata.offset, postingListsMetadata.length);
+ }
+
+ private static List<float[]> cacheOriginalVectors(GraphIndex<float[]> graph)
+ {
+ var view = graph.getView();
+ return IntStream.range(0, graph.size()).mapToObj(view::getVector).collect(Collectors.toList());
+ }
+
+ public long ramBytesUsed()
+ {
+ return graph.ramBytesUsed();
+ }
+
+ public int size()
+ {
+ return graph.size();
+ }
+
+ /**
+ * @return Row IDs associated with the topK vectors near the query
+ */
+ // VSTODO make this return something with a size
Review Comment:
i think this is obsolete
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1362557283
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1080,19 +1139,26 @@ public void processPartition(RowIterator partition, QueryOptions options, Result
private boolean needsPostQueryOrdering()
{
- // We need post-query ordering only for queries with IN on the partition key and an ORDER BY.
- return restrictions.keyIsInRelation() && !parameters.orderings.isEmpty();
+ // We need post-query ordering only for queries with IN on the partition key and an ORDER BY or index restriction reordering
+ return restrictions.keyIsInRelation() && !parameters.orderings.isEmpty() || needIndexOrdering();
+ }
+
+ private boolean needIndexOrdering()
+ {
+ return orderingComparator != null && orderingComparator.indexOrdering();
}
/**
* Orders results when multiple keys are selected (using IN)
*/
- private void orderResults(ResultSet cqlRows)
+ private void orderResults(ResultSet cqlRows, QueryOptions options)
{
if (cqlRows.size() == 0 || !needsPostQueryOrdering())
return;
- Collections.sort(cqlRows.rows, orderingComparator);
+ Comparator<List<ByteBuffer>> comparator = orderingComparator.prepareFor(table, getRowFilter(options), options);
+ if (comparator != null)
+ Collections.sort(cqlRows.rows, comparator);
Review Comment:
Actually, the order isn't necessarily first by the indexed column, then by primary key. The rows gotten from the ANN index don't have that order, and the primary key order is applied after applying the query limit, I think. So we can have surprising cases like this one:
```java
@Test
public void sortLimitTest()
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int, c int, v vector<float, 1>, PRIMARY KEY(k, c))"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}"));
ICoordinator coordinator = cluster.coordinator(1);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (0, 1, [1])"), ConsistencyLevel.ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (0, 2, [1])"), ConsistencyLevel.ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (0, 3, [1])"), ConsistencyLevel.ALL);
Object[][] rows = coordinator.execute(withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [2] LIMIT 10"), ConsistencyLevel.ONE);
assertRows(rows,
row(0, 1, CQLTester.vector(1f)),
row(0, 2, CQLTester.vector(1f)),
row(0, 3, CQLTester.vector(1f)));
rows = coordinator.execute(withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [2] LIMIT 1"), ConsistencyLevel.ONE);
assertRows(rows,
row(0, 1, CQLTester.vector(1f))); // returns [0, 3, 1] instead of [0, 1, 1]!!!
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1364016948
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
Review Comment:
It's not so, much of limitation, as a design decision. Splitting, the index up into N pieces slows down search by (almost) a factor of N.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357008294
##########
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java:
##########
@@ -424,8 +424,9 @@ public enum CassandraRelevantProperties
SAI_INTERSECTION_CLAUSE_LIMIT("cassandra.sai.intersection_clause_limit", "2"),
/** Latest version to be used for SAI index writing */
SAI_LATEST_VERSION("cassandra.sai.latest_version", "aa"),
- SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size_kb", "5"),
- SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size_kb", "1"),
+ SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size", "5KiB"),
+ SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size", "1KiB"),
+ SAI_MAX_VECTOR_TERM_SIZE("cassandra.sai.max_vector_term_size", "16KiB"),
Review Comment:
We're allowing a max of 8k dimension which would be 32KB in Astra, fwiw
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1355010940
##########
src/java/org/apache/cassandra/db/marshal/VectorType.java:
##########
@@ -365,6 +366,7 @@ public abstract class VectorSerializer extends TypeSerializer<List<T>>
public abstract <V> List<V> split(V buffer, ValueAccessor<V> accessor);
public abstract <V> V serializeRaw(List<V> elements, ValueAccessor<V> accessor);
+ public abstract float[] deserializeFloatArray(ByteBuffer input);
Review Comment:
Having looked again, I see that the `deserializeFloatArray` is only implemented for the `Fixed` variant of the serializer, so it can be removed in favour of the `composeAsFloat`. The `deserializeFloatArray` was added because it could shortcut the deserialization by using the `FloatBuffer` to get the array. I have added a `ValueAccessor.toFloatArray` method in order to support this optimisation for `ByteBuffer`s.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "JeremiahDJordan (via GitHub)" <gi...@apache.org>.
JeremiahDJordan commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357060159
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -281,14 +288,29 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
Selectors selectors = selection.newSelectors(options);
AggregationSpecification aggregationSpec = getAggregationSpec(options);
- ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
- nowInSec, userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+ DataLimits limit = getDataLimits(userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+
+ // Handle additional validation for topK queries
+ if (restrictions.isTopK())
+ {
+ checkFalse(aggregationSpec != null, TOPK_AGGREGATION_ERROR);
+
+ checkFalse(options.getConsistency() != ConsistencyLevel.ONE &&
Review Comment:
Rather than failing queries not using LOCAL_ONE/ONE I think we should downgrade the CL to ONE/LOCAL_ONE (depending on if the incoming query used LOCAL_) and set a Client Warn that it happened.
It will be very painful for end users if they have tospecially manage that ANN read queries use ONE when everything else in their app is using LOCAL_QUORUM...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370113494
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/OptimizeFor.java:
##########
@@ -0,0 +1,30 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+public enum OptimizeFor
Review Comment:
This class could use some JavaDoc. Maybe something saying that this is an index config option and that optimizing for recall the ANN index will be asked for more results, etc.
Maybe we could move `VectorIndexSegmentSearcher#topKFor` to this class.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1371640221
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,35 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
+ "They don't yet support SELECT queries with:\n" +
+ " * Consistency level higher than ONE/LOCAL_ONE.\n" +
+ " * Paging.\n" +
+ " * No LIMIT clauses.\n" +
+ " * PER PARTITION LIMIT clauses.\n" +
+ " * GROUP BY clauses.\n" +
+ " * Aggregation functions.\n" +
+ " * Filters on columns without a SAI index.\n" +
+ "Most of these limitations are planned to be removed in future versions.";
Review Comment:
Agreed, we don't need to say this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1350162595
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -109,6 +112,9 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(SelectStatement.logger, 1, TimeUnit.MINUTES);
public static final int DEFAULT_PAGE_SIZE = 10000;
+ public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top-k queries can only be run with consistency level one. Consistency level %s was used.";
Review Comment:
```suggestion
public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top-k queries can only be run with consistency level ONE/LOCAL_ONE. Consistency level %s was used.";
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356733277
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,29 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
Review Comment:
I was thinking that we could format the limitations as a bullet list, and mention aggregate functions and `GROUP BY` instead of aggregation queries, which is kind of internal terminology. Something like:
```java
public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
"They don't yet support SELECT queries with:\n" +
" * Consistency level higher than ONE/LOCAL_ONE.\n" +
" * Paging.\n" +
" * No LIMIT clauses.\n" +
" * PER PARTITION LIMIT clauses.\n" +
" * GROUP BY clauses.\n" +
" * Aggregation functions.\n" +
" * Filters on columns without a SAI index.\n" +
"These limitations are planned to be removed in future versions.";
```
And perhaps we should say that _most of_ these limitations are planned to be removed in future versions, since it's not clear that we will able to get rid of all of them. wdyt?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356959179
##########
src/java/org/apache/cassandra/index/sai/plan/Expression.java:
##########
@@ -172,6 +178,11 @@ public Expression add(Operator op, ByteBuffer value)
*/
public boolean isSatisfiedBy(ByteBuffer columnValue)
{
+ // If the expression represents an ANN ordering then we return true because the actual result
+ // is approximate and with rarely / never match the expression value
Review Comment:
```suggestion
// is approximate and will rarely / never match the expression value
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356796817
##########
test/unit/org/apache/cassandra/inject/injections.md:
##########
@@ -0,0 +1,356 @@
+<!---
+ 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.
+-->
+
+## Injecting hooks dynamically in testing
+
+The testing infrastructure is equipped with Byteman, which allows to
+inject hooks into the running code. In other words, at any point in
+the test, you can tell the particular nodes, to do something at exact
+location in the code - like you would set a conditional breakpoint
+and when the debugger breaks, you can check something or perform some
+additional actions. In short, it is particularly useful to:
+
+- synchronizing the test code and the nodes
+- synchronizing the nodes between each other
+- running actions step by step
+- counting invocations
+- tracing invocations
+- force throwing exceptions
+- and combinations of the above...
+
+Here you can find a short introduction and some examples how can it be
+used.
+
+### Basics
+
+Byteman works on rules which are the units defining what, when, and
+where should be invoked. That is, a single rule includes: the invoke
+point, bindings, run condition and actions.
+
+**Invoke point** is a location in the code where the actions should
+be hooked. Usually it is a class and method, but it can be specified
+more precisely, like - entering method, exiting method, invoking some
+method inside, particular line of code and so on.
+
+**Bindings** are just some constant definitions which can be used then
+in run condition and actions.
+
+**Run condition** is a logical expression which determines whether
+the rule should be invoked.
+
+**Actions** are the statements to be invoked.
+
+In DSE, we have `Rule` class which reflects a single Byteman rule.
Review Comment:
I have deleted that file because it is referencing an old version of the injections that isn't valid for the version being used.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357154655
##########
src/java/org/apache/cassandra/index/sai/memory/VectorMemtableIndex.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * 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.index.sai.memory;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.NavigableSet;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.atomic.LongAdder;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import javax.annotation.Nullable;
+
+import io.github.jbellis.jvector.util.Bits;
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.PartitionPosition;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.VectorQueryContext;
+import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.index.sai.utils.PrimaryKeys;
+import org.apache.cassandra.index.sai.utils.RangeUtil;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.cassandra.utils.Pair;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+
+public class VectorMemtableIndex extends MemoryIndex
Review Comment:
VectorMemoryIndex for consistency?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1363988610
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,35 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
+ "They don't yet support SELECT queries with:\n" +
+ " * Consistency level higher than ONE/LOCAL_ONE.\n" +
+ " * Paging.\n" +
+ " * No LIMIT clauses.\n" +
+ " * PER PARTITION LIMIT clauses.\n" +
+ " * GROUP BY clauses.\n" +
+ " * Aggregation functions.\n" +
+ " * Filters on columns without a SAI index.\n" +
+ "Most of these limitations are planned to be removed in future versions.";
Review Comment:
After recent discussions on the viability of supporting CL>ONE/LOCAL_ONE and paging, maybe we should just remove this line for now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1362397268
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -274,10 +292,54 @@ else if (relation.isLIKE())
throw invalidRequest("Non PRIMARY KEY columns found in where clause: %s ",
Joiner.on(", ").join(nonPrimaryKeyColumns));
}
+
+ var annRestriction = Streams.stream(nonPrimaryKeyRestrictions).filter(SingleRestriction::isANN).findFirst();
+ if (annRestriction.isPresent())
+ {
+ // If there is an ANN restriction then it must be for a vector<float, n> column, and it must have an index
+ var annColumn = annRestriction.get().getFirstColumn();
+
+ if (!annColumn.type.isVector() || !(((VectorType<?>)annColumn.type).elementType instanceof FloatType))
+ throw invalidRequest(StatementRestrictions.ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE);
+ if (indexRegistry == null || indexRegistry.listIndexes().stream().noneMatch(i -> i.dependsOn(annColumn)))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE);
+ // We do not allow ANN queries using partition key restrictions that need filtering
+ if (partitionKeyRestrictions.needFiltering(table))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE);
+ // We do not allow ANN query filtering using non-indexed columns
+ var nonAnnColumns = Streams.stream(nonPrimaryKeyRestrictions).filter(r -> !r.isANN()).map(r -> r.getFirstColumn()).collect(Collectors.toList());
Review Comment:
```suggestion
var nonAnnColumns = Streams.stream(nonPrimaryKeyRestrictions).filter(r -> !r.isANN()).map(Restriction::getFirstColumn).collect(Collectors.toList());
```
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -274,10 +292,54 @@ else if (relation.isLIKE())
throw invalidRequest("Non PRIMARY KEY columns found in where clause: %s ",
Joiner.on(", ").join(nonPrimaryKeyColumns));
}
+
+ var annRestriction = Streams.stream(nonPrimaryKeyRestrictions).filter(SingleRestriction::isANN).findFirst();
+ if (annRestriction.isPresent())
+ {
+ // If there is an ANN restriction then it must be for a vector<float, n> column, and it must have an index
+ var annColumn = annRestriction.get().getFirstColumn();
+
+ if (!annColumn.type.isVector() || !(((VectorType<?>)annColumn.type).elementType instanceof FloatType))
+ throw invalidRequest(StatementRestrictions.ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE);
+ if (indexRegistry == null || indexRegistry.listIndexes().stream().noneMatch(i -> i.dependsOn(annColumn)))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE);
+ // We do not allow ANN queries using partition key restrictions that need filtering
+ if (partitionKeyRestrictions.needFiltering(table))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE);
+ // We do not allow ANN query filtering using non-indexed columns
+ var nonAnnColumns = Streams.stream(nonPrimaryKeyRestrictions).filter(r -> !r.isANN()).map(r -> r.getFirstColumn()).collect(Collectors.toList());
+ var clusteringColumns = clusteringColumnsRestrictions.getColumnDefinitions();
+ if (!nonAnnColumns.isEmpty() || !clusteringColumns.isEmpty())
+ {
+ var nonIndexedColumns = Stream.concat(nonAnnColumns.stream(), clusteringColumns.stream())
+ .filter(c -> indexRegistry == null || indexRegistry.listIndexes()
+ .stream()
+ .noneMatch(i -> i.dependsOn(c)))
+ .collect(Collectors.toList());
+ if (!nonIndexedColumns.isEmpty())
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE);
Review Comment:
If we have the following schema:
```
CREATE TABLE %s.t (k int, c int, v vector<float, 1>, PRIMARY KEY(k, c))
CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex'
```
we can add a non-indexed restriction for the partition key, for example:
```
SELECT * FROM %s.t WHERE k=0 ORDER BY v ANN OF [9] LIMIT 10
```
however, this prevents us from using a clustering prefix, even if that doesn't require `ALLOW FILTERING`:
```
SELECT * FROM %s.t WHERE k=0 AND c>2 ORDER BY v ANN OF [9] LIMIT 10
```
Is this correct? Does the additional clustering require post-filtering?
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -274,10 +292,54 @@ else if (relation.isLIKE())
throw invalidRequest("Non PRIMARY KEY columns found in where clause: %s ",
Joiner.on(", ").join(nonPrimaryKeyColumns));
}
+
+ var annRestriction = Streams.stream(nonPrimaryKeyRestrictions).filter(SingleRestriction::isANN).findFirst();
+ if (annRestriction.isPresent())
+ {
+ // If there is an ANN restriction then it must be for a vector<float, n> column, and it must have an index
+ var annColumn = annRestriction.get().getFirstColumn();
+
+ if (!annColumn.type.isVector() || !(((VectorType<?>)annColumn.type).elementType instanceof FloatType))
+ throw invalidRequest(StatementRestrictions.ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE);
+ if (indexRegistry == null || indexRegistry.listIndexes().stream().noneMatch(i -> i.dependsOn(annColumn)))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE);
+ // We do not allow ANN queries using partition key restrictions that need filtering
+ if (partitionKeyRestrictions.needFiltering(table))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE);
+ // We do not allow ANN query filtering using non-indexed columns
+ var nonAnnColumns = Streams.stream(nonPrimaryKeyRestrictions).filter(r -> !r.isANN()).map(r -> r.getFirstColumn()).collect(Collectors.toList());
+ var clusteringColumns = clusteringColumnsRestrictions.getColumnDefinitions();
+ if (!nonAnnColumns.isEmpty() || !clusteringColumns.isEmpty())
+ {
+ var nonIndexedColumns = Stream.concat(nonAnnColumns.stream(), clusteringColumns.stream())
+ .filter(c -> indexRegistry == null || indexRegistry.listIndexes()
Review Comment:
Nit: at this point `indexRegistry` cannot be null (in that case we throw `ANN_REQUIRES_INDEX_MESSAGE` a few lines above)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#issuecomment-1775435082
@adelapena @jbellis I believe I have now addressed the current review comments. I will carry on with CI getting the tests work reliably. I am nearly there with that.
If you could please have another look, I would be grateful.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1337096020
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -288,7 +291,7 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
query.trackWarnings();
ResultMessage.Rows rows;
- if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize)))
+ if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize) || query.isTopK()))
Review Comment:
I don't think that we should be checking the actual maximum size of the limit here because that would be index specific but I have added a check that a limit has been specified and that it is `<=` to the page size (if the page size is defined).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341419901
##########
src/java/org/apache/cassandra/index/sai/disk/format/IndexComponent.java:
##########
@@ -50,6 +50,11 @@ public enum IndexComponent
*/
TERMS_DATA("TermsData"),
+ /**
+ * Product Quantization store used to store compressed vectors for the vector index
+ */
+ PQ("PQ"),
Review Comment:
That's an alternative but it is describing how the vectors are stored rather than what is being stored in the component. How about we just call it `VECTORS`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341358891
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -664,7 +712,7 @@ public void insertRow(Row row)
@Override
public void updateRow(Row oldRow, Row newRow)
{
- insertRow(newRow);
+ indexContext.getMemtableIndexManager().update(key, oldRow, newRow, memtable);
Review Comment:
Yes we do, I'll add memory management.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341556033
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -276,8 +290,25 @@ else if (relation.isLIKE())
}
if (hasQueriableIndex)
usesSecondaryIndexing = true;
- else if (!allowFiltering && requiresAllowFilteringIfNotSpecified())
- throw invalidRequest(allowFilteringMessage(state));
+ else
+ {
+ if (nonPrimaryKeyRestrictions.hasAnn())
+ {
+ var vectorColumn = nonPrimaryKeyRestrictions.getColumnDefs().stream().filter(c -> c.type.isVector()).findFirst();
Review Comment:
I had to play around with this a bit but the following test fails because of this:
```
@Test
public void multipleVectorColumnsTest()
{
createTable("CREATE TABLE %s (k int PRIMARY KEY, v1 vector<float, 1>, v2 vector<float, 1>)");
createIndex("CREATE CUSTOM INDEX ON %s(v1) USING 'StorageAttachedIndex'");
execute("INSERT INTO %s (k, v1, v2) VALUES (1, [1], [2])");
execute("SELECT * FROM %s WHERE v1 = [1] ORDER BY v2 ANN OF [2] ALLOW FILTERING");
}
```
The problem is that we have a `vectorColumn` and we have a `vectorIndex`, just neither of them in the correct places. We need to check that we have the correct `vectorColumn` for the each test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341335212
##########
src/java/org/apache/cassandra/cql3/SingleColumnRelation.java:
##########
@@ -264,6 +265,16 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
return new SingleColumnRestriction.LikeRestriction(columnDef, operator, term);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ ColumnMetadata columnDef = table.getExistingColumn(entity);
+ if (!(columnDef.type instanceof VectorType))
Review Comment:
I don't see those changes yet, maybe they have to be pushed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341465806
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -266,6 +278,8 @@ else if (relation.isLIKE())
// there is restrictions not covered by the PK.
if (!nonPrimaryKeyRestrictions.isEmpty())
{
+ var columnRestrictions = allColumnRestrictions(clusteringColumnsRestrictions, nonPrimaryKeyRestrictions);
Review Comment:
I have refactored this block a bit to handle the validation a bit better. Where possible I have added extra tests as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1372077193
##########
src/java/org/apache/cassandra/index/sai/plan/vector/VectorQueryController.java:
##########
@@ -0,0 +1,247 @@
+/*
+ * 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.index.sai.plan.vector;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.memtable.Memtable;
+import org.apache.cassandra.exceptions.QueryCancelledException;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.disk.SSTableRowIdKeyRangeIterator;
+import org.apache.cassandra.index.sai.disk.v1.postings.MergePostingList;
+import org.apache.cassandra.index.sai.disk.v1.vector.CheckpointingIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator;
+import org.apache.cassandra.index.sai.metrics.TableQueryMetrics;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.plan.StorageAttachedIndexQueryController;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.postings.PostingListIntersection;
+import org.apache.cassandra.index.sai.postings.PostingListUnion;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.utils.Throwables;
+
+public class VectorQueryController extends StorageAttachedIndexQueryController
+{
+ private static final Logger logger = LoggerFactory.getLogger(VectorQueryController.class);
+
+ public VectorQueryController(ColumnFamilyStore cfs,
+ ReadCommand command,
+ RowFilter filterOperation,
+ QueryContext queryContext,
+ TableQueryMetrics tableQueryMetrics)
+ {
+ super(cfs, command, filterOperation, queryContext, tableQueryMetrics);
+ }
+
+ /**
+ * Build a {@link KeyRangeIterator.Builder} from the given list of expressions by applying given operation (OR/AND).
+ * Building of such builder involves index search, results of which are persisted in the internal resources list
+ *
+ * @param expressions The expressions to build range iterator from (expressions with not results are ignored).
+ *
+ * @return range iterator builder based on given expressions and operation type.
+ */
+ public KeyRangeIterator.Builder getIndexQueryResults(Collection<Expression> expressions)
+ {
+ // FIXME having this at the expression level means that it only gets applied to Nodes at that level;
+ // moving it to ORDER BY will allow us to apply it correctly for other Node sub-trees
+ var annExpressionInHybridSearch = getAnnExpressionInHybridSearch(expressions);
+ boolean isAnnHybridSearch = annExpressionInHybridSearch != null;
+ if (isAnnHybridSearch)
+ expressions = expressions.stream().filter(e -> e != annExpressionInHybridSearch).collect(Collectors.toList());
+
+ // search memtable before referencing sstable indexes; otherwise we may miss newly flushed memtable index
+ var iteratorsByMemtable = expressions.stream()
+ .filter(expr -> !expr.context.isNotIndexed())
+ .flatMap(expr -> expr.context.getMemtableIndexManager().iteratorsForSearch(queryContext, expr, mergeRange).stream())
+ .collect(Collectors.groupingBy(pair -> pair.left, Collectors.mapping(pair -> pair.right, Collectors.toList())));
+
+ VectorQueryViewBuilder.QueryView queryView = new VectorQueryViewBuilder(expressions, mergeRange).build();
+ // in case of ANN query in hybrid search, we have to reference ANN sstable indexes separately because queryView doesn't include ANN sstable indexes
+ VectorQueryViewBuilder.QueryView annQueryViewInHybridSearch = isAnnHybridSearch ? new VectorQueryViewBuilder(Collections.singleton(annExpressionInHybridSearch), mergeRange).build() : null;
+
+ try
+ {
+ List<KeyRangeIterator> sstableIntersections = new ArrayList<>(queryView.view.size());
+ for (Map.Entry<SSTableReader, List<VectorQueryViewBuilder.IndexExpression>> viewEntry : queryView.view.entrySet())
+ {
+ PostingList postings = createRowIdIterator(viewEntry.getValue(), isAnnHybridSearch);
+ if (isAnnHybridSearch)
+ sstableIntersections.add(reorderAndLimitBySSTableRowIds(postings, viewEntry.getKey(), annQueryViewInHybridSearch));
+ else
+ {
+ PrimaryKeyMap.Factory pkFactory = viewEntry.getValue().iterator().next().index.getSSTableContext().primaryKeyMapFactory;
+ sstableIntersections.add(convertToPrimaryKeyIterator(pkFactory, postings));
+ }
+ }
+
+ List<KeyRangeIterator> memtableIntersections = iteratorsByMemtable.entrySet()
+ .stream()
+ .map(e -> {
+ // we need to do all the intersections at the index level, or ordering won't work
+ KeyRangeIterator it = buildIterator(e.getValue(), isAnnHybridSearch);
+ if (isAnnHybridSearch)
+ it = reorderAndLimitBy(it, e.getKey(), annExpressionInHybridSearch);
+ return it;
+ })
+ .collect(Collectors.toList());
+
+ Iterable<KeyRangeIterator> allIntersections = Iterables.concat(sstableIntersections, memtableIntersections);
Review Comment:
I didn't do an exhaustive comparison but it looks like we're still missing the tests that 49a51690 added in ComplexQueryTest, VectorTypeTest, and VectorUpdateDeleteTest.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1368576593
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -109,6 +112,15 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(SelectStatement.logger, 1, TimeUnit.MINUTES);
public static final int DEFAULT_PAGE_SIZE = 10000;
+ public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top-K queries can only be run with consistency level ONE/LOCAL_ONE. Consistency level %s was used.";
+ public static final String TOPK_LIMIT_ERROR = "Top-K queries must have a limit specified and the limit must be less than the query page size";
+ public static final String TOPK_PARTITION_LIMIT_ERROR = "Top-K queries do not support per-partition limits";
+ public static final String TOPK_AGGREGATION_ERROR = "Top-K queries can not be run with aggregation";
+ public static final String TOPK_CONSISTENCY_LEVEL_WARNING = "Top-K queries can only be run with consistency level ONE " +
+ "/ LOCAL_ONE / NODE_LOCAL. Consistency level %s was used. " +
+ "Downgrading to consistency level %s.";
Review Comment:
Ok, I've changed this in the code to check for `NODE_LOCAL` as well and I've changed the warning message to make it clearer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "michaelsembwever (via GitHub)" <gi...@apache.org>.
michaelsembwever closed pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
URL: https://github.com/apache/cassandra/pull/2673
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1371825799
##########
src/java/org/apache/cassandra/index/sai/plan/VectorTopKProcessor.java:
##########
@@ -0,0 +1,198 @@
+/*
+ * 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.index.sai.plan;
+
+import java.nio.ByteBuffer;
+import java.util.Comparator;
+import java.util.Optional;
+import java.util.PriorityQueue;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import javax.annotation.Nullable;
+
+import com.google.common.base.Preconditions;
+import org.apache.commons.lang3.tuple.Triple;
+
+import org.apache.cassandra.cql3.Operator;
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.partitions.BasePartitionIterator;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.rows.BaseRowIterator;
+import org.apache.cassandra.db.rows.Row;
+import org.apache.cassandra.db.rows.Unfiltered;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.index.SecondaryIndexManager;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.index.sai.utils.InMemoryPartitionIterator;
+import org.apache.cassandra.index.sai.utils.InMemoryUnfilteredPartitionIterator;
+import org.apache.cassandra.index.sai.utils.PartitionInfo;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.utils.FBUtilities;
+import org.apache.cassandra.utils.Pair;
+
+/**
+ * Processor that scans all rows from given partitions and selects rows with top-k scores based on vector indexes.
+ * <p>
+ * This processor performs the following steps:
+ * - collect rows with score into PriorityQueue that sorts rows based on score. If there are multiple vector indexes,
+ * the final score is the sum of all vector index scores.
+ * - remove rows with the lowest scores from PQ if PQ size exceeds limit
+ * - return rows from PQ in primary key order to client
+ * <p>
+ * Note that recall will be lower with paging, because:
+ * - page size is used as limit
+ * - for the first query, coordinator returns global top page-size rows within entire ring
+ * - for the subsequent queries, coordinators returns global top page-size rows withom range from last-returned-row to max token
Review Comment:
Yes, I have removed that whole section on paging.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327696483
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -135,6 +135,9 @@ private boolean needsReplicaFilteringProtection()
if (command.rowFilter().isEmpty())
return false;
+ if (command.isTopK())
+ return false;
Review Comment:
Here is a dtest showing how a vector search with CL>ONE can wrongly return a row that has been deleted:
```java
try (Cluster cluster = init(Cluster.build(2).withConfig(config -> config.with(GOSSIP, NETWORK)).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int PRIMARY KEY, v vector<float, 1>)"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex'"));
SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
// insert a row in both nodes, and then remove it in only one of the nodes
cluster.coordinator(1).execute(withKeyspace("INSERT INTO %s.t (k, v) VALUES (1, [1])"), ConsistencyLevel.QUORUM);
cluster.get(1).executeInternal(withKeyspace("DELETE FROM %s.t WHERE k=1"));
String query = withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [1] LIMIT 10");
Object[][] rows = cluster.coordinator(1).execute(query, ConsistencyLevel.QUORUM);
assertThat(rows).isEmpty(); // wrongly returns the deleted row
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1325818624
##########
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java:
##########
@@ -420,12 +420,16 @@ public enum CassandraRelevantProperties
// SAI specific properties
+ /** Whether to allow the user to specify custom options to the hnsw index */
+ SAI_HNSW_ALLOW_CUSTOM_PARAMETERS("cassandra.sai.hnsw.allow_custom_parameters", "false"),
Review Comment:
This property could be with the other vector search properties, a few lines below.
##########
src/java/org/apache/cassandra/cql3/Operator.java:
##########
@@ -258,6 +258,20 @@ public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteB
{
throw new UnsupportedOperationException();
}
+ },
+ ANN(15)
+ {
+ @Override
+ public String toString()
+ {
+ return "ANN";
+ }
+
+ @Override
+ public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteBuffer rightOperand)
+ {
+ return true;
Review Comment:
This probably deserves a comment explaining why it is no-op.
##########
src/java/org/apache/cassandra/cql3/Ordering.java:
##########
@@ -0,0 +1,187 @@
+/*
+ * 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;
+
+import org.apache.cassandra.cql3.restrictions.SingleColumnRestriction;
+import org.apache.cassandra.cql3.restrictions.SingleRestriction;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.TableMetadata;
+
+/**
+ * A single element of an ORDER BY clause.
+ * <code>ORDER BY ordering1 [, ordering2 [, ...]] </code>
+ * <p>
+ * An ordering comprises an expression that produces the values to compare against each other
+ * and a sorting direction (ASC, DESC).
+ */
+public class Ordering
+{
+ public final Expression expression;
+ public final Direction direction;
+
+ public Ordering(Expression expression, Direction direction)
+ {
+ this.expression = expression;
+ this.direction = direction;
+ }
+
+ public interface Expression
+ {
+ default boolean hasNonClusteredOrdering()
+ {
+ return false;
+ }
+
+ default SingleRestriction toRestriction()
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ ColumnMetadata getColumn();
+ }
+
+ /**
+ * Represents a single column in
+ * <code>ORDER BY column</code>
Review Comment:
Nit: could be single-line
```suggestion
* Represents a single column in {@code ORDER BY column}.
```
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -135,6 +135,9 @@ private boolean needsReplicaFilteringProtection()
if (command.rowFilter().isEmpty())
return false;
+ if (command.isTopK())
+ return false;
Review Comment:
This method is called `needsReplicaFilteringProtection`, but it actually is a mix of whether the query needs and supports replica filtering protection. Maybe we could rename it to `useReplicaFilteringProtection`.
I understand that top-k queries with CL_ONE do need replica filtering protection, but we aren't able to provide it, is that right?
If a top-k query with CL>ONE gets a row from one replica and nothing from the other, we cannot know if that row has been changed in the other node. That can lead to the resurrection of deleted/stale rows. Also, read repair might send the stale row to the silent replica. Am I missing something in top-k queries that prevents this situation? I haven't read through all the patch yet so I might be missing something.
I guess a possible approach would be just forbidding top-k queries with CL>ONE.
I understand that ANN in particular, differently to general top-k queries, only has to produce approximate results. So I guess we could use RFP for ANN queries and, if the silent replica gives us something that is not very far away from the index result, just let it pass.
A third approach could be simply failing the queries if RFP actually finds a silent replica. In theory, this shouldn't be a very common situation. Or, instead of failing, it could trigger RR and query the silent replica again. This later approach, if viable, would be quite involved so we could do it in a separate patch.
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,465 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_TABLE_TWO_VECTORS = "CREATE TABLE %%s (pk int primary key, val1 vector<float, %d>, val2 vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final String INVALID_LIMIT_MESSAGE = "Use of ANN OF in an ORDER BY clause requires a LIMIT that is not greater than 1000";
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
+ .withConfig(config -> config.with(GOSSIP).with(NETWORK))
+ .start();
+
+ cluster.schemaChange(withKeyspace(String.format(CREATE_KEYSPACE, RF)));
+ }
+
+ @AfterClass
+ public static void closeCluster()
+ {
+ if (cluster != null)
+ cluster.close();
+ }
+
+ @Before
+ public void before()
+ {
+ table = "table_" + seq.getAndIncrement();
+ dimensionCount = SAITester.getRandom().nextIntBetween(100, 2048);
+ }
+
+ @After
+ public void after()
+ {
+ cluster.schemaChange(formatQuery("DROP TABLE IF EXISTS %s"));
+ }
+
+ @Test
+ public void testVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecall = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithoutLimit(randomVector(), vectorCount))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ int pageSize = SAITester.getRandom().nextIntBetween(40, 70);
+ limit = SAITester.getRandom().nextIntBetween(20, 50);
+ result = searchWithPageAndLimit(queryVector, pageSize, limit);
+
+ resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecallWithPaging = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecallWithPaging).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithPageWithoutLimit(randomVector(), 10))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ // query on-disk index
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ queryVector = randomVector();
+ result = searchWithLimit(queryVector, limit);
+ double sstableRecall = getRecall(vectors, queryVector, getVectors(result));
+ assertThat(sstableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testMultiSSTablesVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+ // disable compaction
+ String tableName = table;
+ cluster.forEach(n -> n.runOnInstance(() -> {
+ Keyspace keyspace = Keyspace.open(KEYSPACE);
+ keyspace.getColumnFamilyStore(tableName).disableAutoCompaction();
+ }));
+
+ int vectorCountPerSSTable = SAITester.getRandom().nextIntBetween(200, 500);
+ int sstableCount = SAITester.getRandom().nextIntBetween(3, 5);
+ List<float[]> allVectors = new ArrayList<>(sstableCount * vectorCountPerSSTable);
+
+ int pk = 0;
+ for (int i = 0; i < sstableCount; i++)
+ {
+ List<float[]> vectors = generateVectors(vectorCountPerSSTable);
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ allVectors.addAll(vectors);
+ cluster.forEach(n -> n.flush(KEYSPACE));
+ }
+
+ // query multiple sstable indexes in multiple node
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(50, 100), allVectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ // expect recall to be at least 0.8
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double recall = getRecall(allVectors, queryVector, getVectors(result));
+ assertThat(recall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testPartitionRestrictedVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, 1, vectors);
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, 1, vectors);
+ }
+ }
+
+ @Test
+ public void rangeRestrictedTest() throws Throwable
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+
+ int pk = 0;
+ Multimap<Long, float[]> vectorsByToken = ArrayListMultimap.create();
+ for (float[] vector : vectors)
+ {
+ vectorsByToken.put(Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(pk)).getLongValue(), vector);
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ',' + vectorString(vector) + " )");
+ }
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index with existing key:
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+ }
+
+ private List<float[]> searchWithRange(float[] queryVector, long minToken, long maxToken, int expectedSize) throws Throwable
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE token(pk) <= " + maxToken + " AND token(pk) >= " + minToken + " ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT 1000");
+ assertThat(result).hasNumberOfRows(expectedSize);
+ return getVectors(result);
+ }
+
+ private Object[][] searchWithLimit(float[] queryVector, int limit)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit);
+ assertThat(result).hasNumberOfRows(limit);
+ return result;
+ }
+
+ private Object[][] searchWithoutLimit(float[] queryVector, int results)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector));
+ assertThat(result).hasNumberOfRows(results);
+ return result;
+ }
+
+
+ private Object[][] searchWithPageWithoutLimit(float[] queryVector, int pageSize)
+ {
+ return executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector), pageSize);
+ }
+
+ private Object[][] searchWithPageAndLimit(float[] queryVector, int pageSize, int limit)
+ {
+ // we don't know how many will be returned in case of paging, because coordinator resumes from last-returned-row's partiton
+ return executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit, pageSize);
+ }
+
+ private void searchByKeyWithoutLimit(int key, float[] queryVector, List<float[]> vectors)
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE pk = " + key + " AND val ann of " + Arrays.toString(queryVector));
+ assertThat(result).hasNumberOfRows(1);
+ float[] output = getVectors(result).get(0);
+ assertThat(output).isEqualTo(vectors.get(key));
+ }
Review Comment:
Unused method.
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,465 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_TABLE_TWO_VECTORS = "CREATE TABLE %%s (pk int primary key, val1 vector<float, %d>, val2 vector<float, %d>)";
Review Comment:
Nit: unused
##########
src/antlr/Lexer.g:
##########
@@ -224,7 +224,8 @@ K_MASKED: M A S K E D;
K_UNMASK: U N M A S K;
K_SELECT_MASKED: S E L E C T '_' M A S K E D;
-K_VECTOR: V E C T O R;
+K_VECTOR: V E C T O R;
+K_ANN_OF: A N N WS+ O F;
Review Comment:
Is there any particular reason for adding `ANN OF` as a single keyword, instead of just adding `ANN` and reusing the existing `OF` keyword?
##########
src/java/org/apache/cassandra/cql3/Ordering.java:
##########
@@ -0,0 +1,187 @@
+/*
+ * 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;
+
+import org.apache.cassandra.cql3.restrictions.SingleColumnRestriction;
+import org.apache.cassandra.cql3.restrictions.SingleRestriction;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.TableMetadata;
+
+/**
+ * A single element of an ORDER BY clause.
+ * <code>ORDER BY ordering1 [, ordering2 [, ...]] </code>
+ * <p>
+ * An ordering comprises an expression that produces the values to compare against each other
+ * and a sorting direction (ASC, DESC).
+ */
+public class Ordering
+{
+ public final Expression expression;
+ public final Direction direction;
+
+ public Ordering(Expression expression, Direction direction)
+ {
+ this.expression = expression;
+ this.direction = direction;
+ }
+
+ public interface Expression
+ {
+ default boolean hasNonClusteredOrdering()
+ {
+ return false;
+ }
+
+ default SingleRestriction toRestriction()
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ ColumnMetadata getColumn();
+ }
+
+ /**
+ * Represents a single column in
+ * <code>ORDER BY column</code>
+ */
+ public static class SingleColumn implements Expression
+ {
+ public final ColumnMetadata column;
+
+ public SingleColumn(ColumnMetadata column)
+ {
+ this.column = column;
+ }
+
+ @Override
+ public ColumnMetadata getColumn()
+ {
+ return column;
+ }
+ }
+
+ /**
+ * An expression used in Approximate Nearest Neighbor ordering.
+ * <code>ORDER BY column ANN OF value</code>
+ */
+ public static class Ann implements Expression
+ {
+ final ColumnMetadata column;
+ final Term vectorValue;
+
+ public Ann(ColumnMetadata column, Term vectorValue)
+ {
+ this.column = column;
+ this.vectorValue = vectorValue;
+ }
+
+ @Override
+ public boolean hasNonClusteredOrdering()
+ {
+ return true;
+ }
+
+ @Override
+ public SingleRestriction toRestriction()
+ {
+ return new SingleColumnRestriction.AnnRestriction(column, vectorValue);
+ }
+
+ @Override
+ public ColumnMetadata getColumn()
+ {
+ return column;
+ }
+ }
+
+ public enum Direction
+ {ASC, DESC}
+
+
+ /**
+ * Represents the AST of a single element in the ORDER BY clause.
Review Comment:
What is AST?
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,465 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_TABLE_TWO_VECTORS = "CREATE TABLE %%s (pk int primary key, val1 vector<float, %d>, val2 vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final String INVALID_LIMIT_MESSAGE = "Use of ANN OF in an ORDER BY clause requires a LIMIT that is not greater than 1000";
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
+ .withConfig(config -> config.with(GOSSIP).with(NETWORK))
+ .start();
+
+ cluster.schemaChange(withKeyspace(String.format(CREATE_KEYSPACE, RF)));
+ }
+
+ @AfterClass
+ public static void closeCluster()
+ {
+ if (cluster != null)
+ cluster.close();
+ }
+
+ @Before
+ public void before()
+ {
+ table = "table_" + seq.getAndIncrement();
+ dimensionCount = SAITester.getRandom().nextIntBetween(100, 2048);
+ }
+
+ @After
+ public void after()
+ {
+ cluster.schemaChange(formatQuery("DROP TABLE IF EXISTS %s"));
+ }
+
+ @Test
+ public void testVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecall = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithoutLimit(randomVector(), vectorCount))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ int pageSize = SAITester.getRandom().nextIntBetween(40, 70);
+ limit = SAITester.getRandom().nextIntBetween(20, 50);
+ result = searchWithPageAndLimit(queryVector, pageSize, limit);
+
+ resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecallWithPaging = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecallWithPaging).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithPageWithoutLimit(randomVector(), 10))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ // query on-disk index
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ queryVector = randomVector();
+ result = searchWithLimit(queryVector, limit);
+ double sstableRecall = getRecall(vectors, queryVector, getVectors(result));
+ assertThat(sstableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testMultiSSTablesVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+ // disable compaction
+ String tableName = table;
+ cluster.forEach(n -> n.runOnInstance(() -> {
+ Keyspace keyspace = Keyspace.open(KEYSPACE);
+ keyspace.getColumnFamilyStore(tableName).disableAutoCompaction();
+ }));
+
+ int vectorCountPerSSTable = SAITester.getRandom().nextIntBetween(200, 500);
+ int sstableCount = SAITester.getRandom().nextIntBetween(3, 5);
+ List<float[]> allVectors = new ArrayList<>(sstableCount * vectorCountPerSSTable);
+
+ int pk = 0;
+ for (int i = 0; i < sstableCount; i++)
+ {
+ List<float[]> vectors = generateVectors(vectorCountPerSSTable);
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ allVectors.addAll(vectors);
+ cluster.forEach(n -> n.flush(KEYSPACE));
+ }
+
+ // query multiple sstable indexes in multiple node
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(50, 100), allVectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ // expect recall to be at least 0.8
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double recall = getRecall(allVectors, queryVector, getVectors(result));
+ assertThat(recall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testPartitionRestrictedVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, 1, vectors);
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, 1, vectors);
+ }
+ }
+
+ @Test
+ public void rangeRestrictedTest() throws Throwable
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+
+ int pk = 0;
+ Multimap<Long, float[]> vectorsByToken = ArrayListMultimap.create();
+ for (float[] vector : vectors)
+ {
+ vectorsByToken.put(Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(pk)).getLongValue(), vector);
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ',' + vectorString(vector) + " )");
+ }
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index with existing key:
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+ }
+
+ private List<float[]> searchWithRange(float[] queryVector, long minToken, long maxToken, int expectedSize) throws Throwable
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE token(pk) <= " + maxToken + " AND token(pk) >= " + minToken + " ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT 1000");
+ assertThat(result).hasNumberOfRows(expectedSize);
+ return getVectors(result);
+ }
+
+ private Object[][] searchWithLimit(float[] queryVector, int limit)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit);
+ assertThat(result).hasNumberOfRows(limit);
+ return result;
+ }
+
+ private Object[][] searchWithoutLimit(float[] queryVector, int results)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector));
+ assertThat(result).hasNumberOfRows(results);
+ return result;
+ }
+
+
+ private Object[][] searchWithPageWithoutLimit(float[] queryVector, int pageSize)
+ {
+ return executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector), pageSize);
+ }
+
+ private Object[][] searchWithPageAndLimit(float[] queryVector, int pageSize, int limit)
+ {
+ // we don't know how many will be returned in case of paging, because coordinator resumes from last-returned-row's partiton
+ return executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit, pageSize);
+ }
+
+ private void searchByKeyWithoutLimit(int key, float[] queryVector, List<float[]> vectors)
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE pk = " + key + " AND val ann of " + Arrays.toString(queryVector));
+ assertThat(result).hasNumberOfRows(1);
+ float[] output = getVectors(result).get(0);
+ assertThat(output).isEqualTo(vectors.get(key));
+ }
+
+ private void searchByKeyWithLimit(int key, float[] queryVector, int limit, List<float[]> vectors)
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE pk = " + key + " ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit);
+ assertThat(result).hasNumberOfRows(1);
+ float[] output = getVectors(result).get(0);
+ assertThat(output).isEqualTo(vectors.get(key));
+ }
+
+ private void assertDescendingScore(float[] queryVector, List<float[]> resultVectors)
+ {
+ float prevScore = -1;
+ for (float[] current : resultVectors)
+ {
+ float score = function.compare(current, queryVector);
+ if (prevScore >= 0)
+ assertThat(score).isLessThanOrEqualTo(prevScore);
+
+ prevScore = score;
+ }
+ }
+
+ private double getRecall(List<float[]> vectors, float[] query, List<float[]> result)
+ {
+ List<float[]> sortedVectors = new ArrayList<>(vectors);
+ sortedVectors.sort((a, b) -> Double.compare(function.compare(b, query), function.compare(a, query)));
+
+ assertThat(sortedVectors).containsAll(result);
+
+ List<float[]> nearestNeighbors = sortedVectors.subList(0, result.size());
+
+ int matches = 0;
+ for (float[] in : nearestNeighbors)
+ {
+ for (float[] out : result)
+ {
+ if (Arrays.compare(in, out) ==0)
+ {
+ matches++;
+ break;
+ }
+ }
+ }
+
+ return matches * 1.0 / result.size();
+ }
+
+ private List<float[]> generateVectors(int vectorCount)
+ {
+ return IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+ }
+
+ private List<float[]> getVectors(Object[][] result)
+ {
+ List<float[]> vectors = new ArrayList<>();
+
+ // verify results are part of inserted vectors
+ for (Object[] obj: result)
+ {
+ List<Float> list = (List<Float>) obj[0];
+ float[] array = new float[list.size()];
+ for (int index = 0; index < list.size(); index++)
+ array[index] = list.get(index);
+ vectors.add(array);
+ }
+
+ return vectors;
+ }
+
+ private String vectorString(float[] vector)
+ {
+ return Arrays.toString(vector);
+ }
+
+ private String randomVectorString()
+ {
+ return vectorString(randomVector());
+ }
Review Comment:
Unused method
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,465 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_TABLE_TWO_VECTORS = "CREATE TABLE %%s (pk int primary key, val1 vector<float, %d>, val2 vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final String INVALID_LIMIT_MESSAGE = "Use of ANN OF in an ORDER BY clause requires a LIMIT that is not greater than 1000";
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
+ .withConfig(config -> config.with(GOSSIP).with(NETWORK))
+ .start();
+
+ cluster.schemaChange(withKeyspace(String.format(CREATE_KEYSPACE, RF)));
+ }
+
+ @AfterClass
+ public static void closeCluster()
+ {
+ if (cluster != null)
+ cluster.close();
+ }
+
+ @Before
+ public void before()
+ {
+ table = "table_" + seq.getAndIncrement();
+ dimensionCount = SAITester.getRandom().nextIntBetween(100, 2048);
+ }
+
+ @After
+ public void after()
+ {
+ cluster.schemaChange(formatQuery("DROP TABLE IF EXISTS %s"));
+ }
+
+ @Test
+ public void testVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecall = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithoutLimit(randomVector(), vectorCount))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ int pageSize = SAITester.getRandom().nextIntBetween(40, 70);
+ limit = SAITester.getRandom().nextIntBetween(20, 50);
+ result = searchWithPageAndLimit(queryVector, pageSize, limit);
+
+ resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecallWithPaging = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecallWithPaging).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithPageWithoutLimit(randomVector(), 10))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ // query on-disk index
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ queryVector = randomVector();
+ result = searchWithLimit(queryVector, limit);
+ double sstableRecall = getRecall(vectors, queryVector, getVectors(result));
+ assertThat(sstableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testMultiSSTablesVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+ // disable compaction
+ String tableName = table;
+ cluster.forEach(n -> n.runOnInstance(() -> {
+ Keyspace keyspace = Keyspace.open(KEYSPACE);
+ keyspace.getColumnFamilyStore(tableName).disableAutoCompaction();
+ }));
+
+ int vectorCountPerSSTable = SAITester.getRandom().nextIntBetween(200, 500);
+ int sstableCount = SAITester.getRandom().nextIntBetween(3, 5);
+ List<float[]> allVectors = new ArrayList<>(sstableCount * vectorCountPerSSTable);
+
+ int pk = 0;
+ for (int i = 0; i < sstableCount; i++)
+ {
+ List<float[]> vectors = generateVectors(vectorCountPerSSTable);
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ allVectors.addAll(vectors);
+ cluster.forEach(n -> n.flush(KEYSPACE));
+ }
+
+ // query multiple sstable indexes in multiple node
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(50, 100), allVectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ // expect recall to be at least 0.8
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double recall = getRecall(allVectors, queryVector, getVectors(result));
+ assertThat(recall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testPartitionRestrictedVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, 1, vectors);
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, 1, vectors);
+ }
+ }
+
+ @Test
+ public void rangeRestrictedTest() throws Throwable
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+
+ int pk = 0;
+ Multimap<Long, float[]> vectorsByToken = ArrayListMultimap.create();
+ for (float[] vector : vectors)
+ {
+ vectorsByToken.put(Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(pk)).getLongValue(), vector);
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ',' + vectorString(vector) + " )");
+ }
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index with existing key:
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+ }
+
+ private List<float[]> searchWithRange(float[] queryVector, long minToken, long maxToken, int expectedSize) throws Throwable
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE token(pk) <= " + maxToken + " AND token(pk) >= " + minToken + " ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT 1000");
+ assertThat(result).hasNumberOfRows(expectedSize);
+ return getVectors(result);
+ }
+
+ private Object[][] searchWithLimit(float[] queryVector, int limit)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit);
+ assertThat(result).hasNumberOfRows(limit);
+ return result;
+ }
+
+ private Object[][] searchWithoutLimit(float[] queryVector, int results)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector));
+ assertThat(result).hasNumberOfRows(results);
+ return result;
+ }
+
+
+ private Object[][] searchWithPageWithoutLimit(float[] queryVector, int pageSize)
+ {
+ return executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector), pageSize);
+ }
+
+ private Object[][] searchWithPageAndLimit(float[] queryVector, int pageSize, int limit)
+ {
+ // we don't know how many will be returned in case of paging, because coordinator resumes from last-returned-row's partiton
+ return executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit, pageSize);
+ }
+
+ private void searchByKeyWithoutLimit(int key, float[] queryVector, List<float[]> vectors)
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE pk = " + key + " AND val ann of " + Arrays.toString(queryVector));
+ assertThat(result).hasNumberOfRows(1);
+ float[] output = getVectors(result).get(0);
+ assertThat(output).isEqualTo(vectors.get(key));
+ }
+
+ private void searchByKeyWithLimit(int key, float[] queryVector, int limit, List<float[]> vectors)
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE pk = " + key + " ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit);
+ assertThat(result).hasNumberOfRows(1);
+ float[] output = getVectors(result).get(0);
+ assertThat(output).isEqualTo(vectors.get(key));
+ }
+
+ private void assertDescendingScore(float[] queryVector, List<float[]> resultVectors)
+ {
+ float prevScore = -1;
+ for (float[] current : resultVectors)
+ {
+ float score = function.compare(current, queryVector);
+ if (prevScore >= 0)
+ assertThat(score).isLessThanOrEqualTo(prevScore);
+
+ prevScore = score;
+ }
+ }
+
+ private double getRecall(List<float[]> vectors, float[] query, List<float[]> result)
+ {
+ List<float[]> sortedVectors = new ArrayList<>(vectors);
+ sortedVectors.sort((a, b) -> Double.compare(function.compare(b, query), function.compare(a, query)));
+
+ assertThat(sortedVectors).containsAll(result);
+
+ List<float[]> nearestNeighbors = sortedVectors.subList(0, result.size());
+
+ int matches = 0;
+ for (float[] in : nearestNeighbors)
+ {
+ for (float[] out : result)
+ {
+ if (Arrays.compare(in, out) ==0)
+ {
+ matches++;
+ break;
+ }
+ }
+ }
+
+ return matches * 1.0 / result.size();
+ }
+
+ private List<float[]> generateVectors(int vectorCount)
+ {
+ return IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+ }
+
+ private List<float[]> getVectors(Object[][] result)
+ {
+ List<float[]> vectors = new ArrayList<>();
+
+ // verify results are part of inserted vectors
+ for (Object[] obj: result)
+ {
+ List<Float> list = (List<Float>) obj[0];
+ float[] array = new float[list.size()];
+ for (int index = 0; index < list.size(); index++)
+ array[index] = list.get(index);
+ vectors.add(array);
+ }
+
+ return vectors;
+ }
+
+ private String vectorString(float[] vector)
+ {
+ return Arrays.toString(vector);
+ }
+
+ private String randomVectorString()
+ {
+ return vectorString(randomVector());
+ }
+
+ private float[] randomVector()
+ {
+ float[] rawVector = new float[dimensionCount];
+ for (int i = 0; i < dimensionCount; i++)
+ {
+ rawVector[i] = SAITester.getRandom().nextFloat();
+ }
+ return rawVector;
+ }
+
+ private static Object[][] execute(String query)
+ {
+ return execute(query, ConsistencyLevel.QUORUM);
+ }
+
+ private static Object[][] executeAll(String query)
+ {
+ return execute(query, ConsistencyLevel.ALL);
+ }
Review Comment:
Unused method
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1333462103
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -40,16 +41,21 @@
import java.util.stream.Collectors;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.Futures; // checkstyle: permit this import
import com.google.common.util.concurrent.ListenableFuture; // checkstyle: permit this import
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
Review Comment:
CQL similarity functions (`org.apache.cassandra.cql3.functions.VectorFcts` and `org.apache.cassandra.cql3.functions.VectorFctsTest`) still import `org.apache.lucene.index.VectorSimilarityFunction`. They should be updated to import `io.github.jbellis.jvector.vector.VectorSimilarityFunction`. They also contain some vars named `luceneFunction` that should be renamed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1333417858
##########
src/java/org/apache/cassandra/db/ReadQuery.java:
##########
@@ -256,4 +256,12 @@ default void maybeValidateIndex()
default void trackWarnings()
{
}
+
+ /**
+ * @return true given read query is a top-k request
Review Comment:
Maybe we could add in the JavaDoc for this method that top-k queries are necessarily based on the ordering provided by an index. Otherwise, the implementation of this method in `ReadCommand` should be based on the `RowFilter` expressions rather than delegating to the index.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341316670
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -109,6 +112,7 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(SelectStatement.logger, 1, TimeUnit.MINUTES);
public static final int DEFAULT_PAGE_SIZE = 10000;
+ public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top K queries can only be run with consistency level one";
Review Comment:
Nit: I think it's more common to use `Top-k` rather than `Top K`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342851973
##########
src/java/org/apache/cassandra/index/sai/IndexContext.java:
##########
@@ -354,6 +371,11 @@ public boolean isLiteral()
return TypeUtil.isLiteral(getValidator());
}
+ public boolean isVector()
+ {
+ return getValidator().isVector() && ((VectorType<?>)getValidator()).elementType instanceof FloatType;
Review Comment:
I'm in favour of 1 for the time being. I don't think treating them as literals makes sense at all. If we were going to index them, using an index other than ANN, it's more likely that we would follow the collections path and effectively tokenize them, or something new.
I'll add some validation to block this case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357027477
##########
src/java/org/apache/cassandra/cql3/restrictions/SingleColumnRestriction.java:
##########
@@ -804,4 +824,70 @@ else if (ByteBufferUtil.startsWith(value, LIKE_WILDCARD))
return Pair.create(operator, newValue);
}
}
+ public static final class AnnRestriction extends SingleColumnRestriction
+ {
+ private final Term value;
+
+ public AnnRestriction(ColumnMetadata columnDef, Term value)
+ {
+ super(columnDef);
+ this.value = value;
+ }
+
+ public ByteBuffer value(QueryOptions options)
+ {
+ return value.bindAndGet(options);
+ }
+
+ @Override
+ public boolean isANN()
+ {
+ return true;
+ }
+
+ @Override
+ public void addFunctionsTo(List<Function> functions)
+ {
+ value.addFunctionsTo(functions);
+ }
+
+ @Override
+ MultiColumnRestriction toMultiColumnRestriction()
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void addToRowFilter(RowFilter filter,
+ IndexRegistry indexRegistry,
+ QueryOptions options)
+ {
+ filter.add(columnDef, Operator.ANN, value.bindAndGet(options));
+ }
+
+ @Override
+ public MultiCBuilder appendTo(MultiCBuilder builder, QueryOptions options)
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public String toString()
+ {
+ return String.format("ANN(%s)", value);
+ }
+
+ @Override
+ public SingleRestriction doMergeWith(SingleRestriction otherRestriction)
+ {
+ // VSTODO not sure if this is right
Review Comment:
obsolete?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356711170
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorSegmentationTest.java:
##########
@@ -0,0 +1,142 @@
+/*
+ * 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.index.sai.cql;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.Test;
+
+import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.db.marshal.FloatType;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentBuilder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class VectorSegmentationTest extends VectorTester
+{
+ private static final int dimension = 100;
+
+ @Test
+ public void testMultipleSegmentsForCreatingIndex() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+
+ int vectorCount = 100;
+ List<float[]> vectors = new ArrayList<>();
+ for (int row = 0; row < vectorCount; row++)
+ {
+ float[] vector = nextVector();
+ vectors.add(vector);
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", row, vector(vector));
+ }
+
+ flush();
+
+ SegmentBuilder.updateLastValidSegmentRowId(17); // 17 rows per segment
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ int limit = 35;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ @Test
+ public void testMultipleSegmentsForCompaction() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ List<float[]> vectors = new ArrayList<>();
+ int rowsPerSSTable = 10;
+ int sstables = 5;
+ int pk = 0;
+ for (int i = 0; i < sstables; i++)
+ {
+ for (int row = 0; row < rowsPerSSTable; row++)
+ {
+ float[] vector = nextVector();
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", pk++, vector(vector));
+ vectors.add(vector);
+ }
+
+ flush();
+ }
+
+ int limit = 30;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+
+
+ SegmentBuilder.updateLastValidSegmentRowId(11); // 11 rows per segment
+ compact();
+
+ queryVector = nextVector();
+ resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ resultVectors = getVectorsFromResult(resultSet);
+ recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ protected Vector<Float> vector(float[] values)
+ {
+ Float[] floats = new Float[values.length];
+ for (int i = 0; i < values.length; i++)
+ floats[i] = values[i];
+
+ return new Vector<>(floats);
+ }
Review Comment:
I think the method is commented out, instead of just removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356978871
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorInvalidQueryTest.java:
##########
@@ -0,0 +1,289 @@
+/*
+ * 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.index.sai.cql;
+
+import com.datastax.driver.core.ResultSet;
+import com.datastax.driver.core.exceptions.InvalidQueryException;
+import org.apache.cassandra.cql3.restrictions.StatementRestrictions;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.Assert.assertEquals;
+
+public class VectorInvalidQueryTest extends SAITester
+{
+ @Test
+ public void cannotCreateEmptyVectorColumn()
+ {
+ assertThatThrownBy(() -> execute(String.format("CREATE TABLE %s.%s (pk int, str_val text, val vector<float, 0>, PRIMARY KEY(pk))",
+ KEYSPACE, createTableName())))
+ .isInstanceOf(InvalidRequestException.class)
+ .hasMessage("vectors may only have positive dimensions; given 0");
+ }
+
+ @Test
+ public void cannotQueryEmptyVectorColumn()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ assertThatThrownBy(() -> execute("SELECT similarity_cosine((vector<float, 0>) [], []) FROM %s"))
+ .isInstanceOf(InvalidRequestException.class)
+ .hasMessage("vectors may only have positive dimensions; given 0");
+ }
+
+ @Test
+ public void cannotInsertWrongNumberOfDimensions()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+
+ assertThatThrownBy(() -> execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0])"))
+ .isInstanceOf(InvalidRequestException.class)
+ .hasMessage("Invalid vector literal for val of type vector<float, 3>; expected 3 elements, but given 2");
+ }
+
+ @Test
+ public void cannotQueryWrongNumberOfDimensions()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ assertThatThrownBy(() -> execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5] LIMIT 5"))
+ .isInstanceOf(InvalidRequestException.class)
+ .hasMessage("Invalid vector literal for val of type vector<float, 3>; expected 3 elements, but given 2");
+ }
+
+ @Test
+ public void testMultiVectorOrderingsNotAllowed() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val1 vector<float, 3>, val2 vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val1) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val2) USING 'StorageAttachedIndex'");
+
+ assertInvalidMessage("Cannot specify more than one ANN ordering",
+ "SELECT * FROM %s ORDER BY val1 ann of [2.5, 3.5, 4.5], val2 ann of [2.1, 3.2, 4.0] LIMIT 2");
+ }
+
+ @Test
+ public void testDescendingVectorOrderingIsNotAllowed() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ assertInvalidMessage("Descending ANN ordering is not supported",
+ "SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] DESC LIMIT 2");
+ }
+
+ @Test
+ public void testVectorOrderingIsNotAllowedWithClusteringOrdering() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, ck int, val vector<float, 3>, PRIMARY KEY(pk, ck))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ assertInvalidMessage("ANN ordering does not support secondary ordering",
+ "SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5], ck ASC LIMIT 2");
+ }
+
+ @Test
+ public void testVectorOrderingIsNotAllowedWithoutIndex() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE,
+ "SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE,
+ "SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5 ALLOW FILTERING");
+ }
+
+ @Test
+ public void testInvalidColumnNameWithAnn() throws Throwable
+ {
+ String table = createTable(KEYSPACE, "CREATE TABLE %s (k int, c int, v int, primary key (k, c))");
+ assertInvalidMessage(String.format("Undefined column name bad_col in table %s", KEYSPACE + '.' + table),
+ "SELECT k from %s ORDER BY bad_col ANN OF [1.0] LIMIT 1");
+ }
+
+ @Test
+ public void cannotPerformNonANNQueryOnVectorIndex() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, ck int, val vector<float, 3>, PRIMARY KEY(pk, ck))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ assertInvalidMessage(StatementRestrictions.VECTOR_INDEXES_ANN_ONLY_MESSAGE,
+ "SELECT * FROM %s WHERE val = [1.0, 2.0, 3.0]");
+ }
+
+ @Test
+ public void cannotOrderWithAnnOnNonVectorColumn() throws Throwable
+ {
+ createTable("CREATE TABLE %s (k int, v int, primary key(k))");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ assertInvalidMessage(StatementRestrictions.ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE,
+ "SELECT * FROM %s ORDER BY v ANN OF 1 LIMIT 1");
+ }
+
+ @Test
+ public void disallowZeroVectorsWithCosineSimilarity()
+ {
+ createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, value vector<float, 2>)");
+ createIndex("CREATE CUSTOM INDEX ON %s(value) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'cosine'}");
+
+ assertThatThrownBy(() -> execute("INSERT INTO %s (pk, value) VALUES (0, [0.0, 0.0])")).isInstanceOf(InvalidRequestException.class);
+ assertThatThrownBy(() -> execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(0.0f, 0.0f))).isInstanceOf(InvalidRequestException.class);
+ assertThatThrownBy(() -> execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(1.0f, Float.NaN))).isInstanceOf(InvalidRequestException.class);
+ assertThatThrownBy(() -> execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(1.0f, Float.POSITIVE_INFINITY))).isInstanceOf(InvalidRequestException.class);
+ assertThatThrownBy(() -> execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(Float.NEGATIVE_INFINITY, 1.0f))).isInstanceOf(InvalidRequestException.class);
+ assertThatThrownBy(() -> execute("SELECT * FROM %s ORDER BY value ann of [0.0, 0.0] LIMIT 2")).isInstanceOf(InvalidRequestException.class);
+ }
+
+ @Test
+ public void mustHaveLimitSpecifiedAndWithinMaxAllowed()
+ {
+ createTable("CREATE TABLE %s (k int PRIMARY KEY, v vector<float, 1>)");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ assertThatThrownBy(() -> executeNet("SELECT * FROM %s WHERE k = 1 ORDER BY v ANN OF [0]"))
+ .isInstanceOf(InvalidQueryException.class).hasMessage(SelectStatement.TOPK_LIMIT_ERROR);
+
+ assertThatThrownBy(() -> executeNet("SELECT * FROM %s WHERE k = 1 ORDER BY v ANN OF [0] LIMIT 1001"))
+ .isInstanceOf(InvalidQueryException.class).hasMessage(String.format(StorageAttachedIndex.ANN_LIMIT_ERROR, IndexWriterConfig.MAX_TOP_K, 1001));
+ }
+
+ @Test
+ public void mustHaveLimitWithinPageSize()
+ {
+ createTable("CREATE TABLE %s (k int PRIMARY KEY, v vector<float, 3>)");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (k, v) VALUES (1, [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (k, v) VALUES (2, [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (k, v) VALUES (3, [3.0, 4.0, 5.0])");
+
+ assertThatThrownBy(() -> executeNetWithPaging("SELECT * FROM %s WHERE k = 1 ORDER BY v ANN OF [2.0, 3.0, 4.0] LIMIT 10", 9))
+ .isInstanceOf(InvalidQueryException.class).hasMessageContaining(SelectStatement.TOPK_LIMIT_ERROR);
+
+ ResultSet resultSet = executeNetWithPaging("SELECT * FROM %s WHERE k = 1 ORDER BY v ANN OF [2.0, 3.0, 4.0] LIMIT 10", 11);
+
+ assertEquals(1, resultSet.all().size());
+ }
+
+ @Test
+ public void cannotHaveAggregationOnANNQuery()
+ {
+ createTable("CREATE TABLE %s (k int PRIMARY KEY, v vector<float, 1>, c int)");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (k, v, c) VALUES (1, [4], 1)");
+ execute("INSERT INTO %s (k, v, c) VALUES (2, [3], 10)");
+ execute("INSERT INTO %s (k, v, c) VALUES (3, [2], 100)");
+ execute("INSERT INTO %s (k, v, c) VALUES (4, [1], 1000)");
+
+ assertThatThrownBy(() -> executeNet("SELECT sum(c) FROM %s WHERE k = 1 ORDER BY v ANN OF [0] LIMIT 4"))
+ .isInstanceOf(InvalidQueryException.class).hasMessage(SelectStatement.TOPK_AGGREGATION_ERROR);
+ }
+
+ @Test
+ public void multipleVectorColumnsInQueryFailCorrectlyTest() throws Throwable
+ {
+ createTable("CREATE TABLE %s (k int PRIMARY KEY, v1 vector<float, 1>, v2 vector<float, 1>)");
+ createIndex("CREATE CUSTOM INDEX ON %s(v1) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (k, v1, v2) VALUES (1, [1], [2])");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE,
+ "SELECT * FROM %s WHERE v1 = [1] ORDER BY v2 ANN OF [2] ALLOW FILTERING");
+ }
+
+ @Test
+ public void annOrderingIsNotAllowedWithoutIndexWhereIndexedColumnExistsInQueryTest() throws Throwable
+ {
+ createTable("CREATE TABLE %s (k int PRIMARY KEY, v vector<float, 1>, c int)");
+ createIndex("CREATE CUSTOM INDEX ON %s(c) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (k, v, c) VALUES (1, [4], 1)");
+ execute("INSERT INTO %s (k, v, c) VALUES (2, [3], 10)");
+ execute("INSERT INTO %s (k, v, c) VALUES (3, [2], 100)");
+ execute("INSERT INTO %s (k, v, c) VALUES (4, [1], 1000)");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE,
+ "SELECT * FROM %s WHERE c >= 100 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
+ }
+
+ @Test
+ public void cannotPostFilterOnNonIndexedColumnWithAnnOrdering() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk1 int, pk2 int, ck1 int, ck2 int, v vector<float, 1>, c int, primary key ((pk1, pk2), ck1, ck2))");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk1, pk2, ck1, ck2, v, c) VALUES (1, 1, 1, 1, [4], 1)");
+ execute("INSERT INTO %s (pk1, pk2, ck1, ck2, v, c) VALUES (2, 2, 1, 1, [3], 10)");
+ execute("INSERT INTO %s (pk1, pk2, ck1, ck2, v, c) VALUES (3, 3, 1, 1, [2], 100)");
+ execute("INSERT INTO %s (pk1, pk2, ck1, ck2, v, c) VALUES (4, 4, 1, 1, [1], 1000)");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE,
+ "SELECT * FROM %s WHERE c >= 100 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE,
+ "SELECT * FROM %s WHERE ck1 >= 0 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE,
+ "SELECT * FROM %s WHERE ck2 = 1 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE,
+ "SELECT * FROM %s WHERE pk1 = 1 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
+
+ assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE,
+ "SELECT * FROM %s WHERE pk2 = 1 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
+ }
Review Comment:
We could add a case for a complete partition key restriction and filtering clustering restrictions (it works as expected):
```java
assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE,
"SELECT * FROM %s WHERE pk1 = 1 AND pk2 = 1 AND ck2 = 1 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
assertInvalidMessage(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE,
"SELECT * FROM %s WHERE token(pk1, pk2) = token(1, 1) AND ck2 = 1 ORDER BY v ANN OF [1] LIMIT 4 ALLOW FILTERING");
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357105248
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java:
##########
@@ -0,0 +1,56 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+
+/**
+ * There are two steps in ordering:
Review Comment:
this explanation is outdated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357046442
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -109,6 +112,10 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(SelectStatement.logger, 1, TimeUnit.MINUTES);
public static final int DEFAULT_PAGE_SIZE = 10000;
+ public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top-k queries can only be run with consistency level ONE/LOCAL_ONE. Consistency level %s was used.";
Review Comment:
Nit: standardize on Top-K (capital K) ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357156519
##########
src/java/org/apache/cassandra/index/sai/memory/VectorMemtableIndex.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * 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.index.sai.memory;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.NavigableSet;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.atomic.LongAdder;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import javax.annotation.Nullable;
+
+import io.github.jbellis.jvector.util.Bits;
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.PartitionPosition;
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.VectorQueryContext;
+import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.index.sai.utils.PrimaryKeys;
+import org.apache.cassandra.index.sai.utils.RangeUtil;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.cassandra.utils.Pair;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+
+public class VectorMemtableIndex extends MemoryIndex
+{
+ private final OnHeapGraph<PrimaryKey> graph;
+ private final LongAdder writeCount = new LongAdder();
+
+ private PrimaryKey minimumKey;
+ private PrimaryKey maximumKey;
+
+ private final NavigableSet<PrimaryKey> primaryKeys = new ConcurrentSkipListSet<>();
+
+ public VectorMemtableIndex(IndexContext indexContext)
+ {
+ super(indexContext);
+ this.graph = new OnHeapGraph<>(indexContext.getValidator(), indexContext.getIndexWriterConfig());
+ }
+
+ @Override
+ public synchronized long add(DecoratedKey key, Clustering<?> clustering, ByteBuffer value)
+ {
+ if (value == null || value.remaining() == 0)
+ return 0;
+
+ var primaryKey = indexContext.hasClustering() ? indexContext.keyFactory().create(key, clustering)
+ : indexContext.keyFactory().create(key);
+ return index(primaryKey, value);
+ }
+
+ private long index(PrimaryKey primaryKey, ByteBuffer value)
+ {
+ updateKeyBounds(primaryKey);
+
+ writeCount.increment();
+ primaryKeys.add(primaryKey);
+ return graph.add(value, primaryKey, OnHeapGraph.InvalidVectorBehavior.FAIL);
+ }
+
+ @Override
+ public long update(DecoratedKey key, Clustering<?> clustering, ByteBuffer oldValue, ByteBuffer newValue)
+ {
+ int oldRemaining = oldValue == null ? 0 : oldValue.remaining();
+ int newRemaining = newValue == null ? 0 : newValue.remaining();
+ if (oldRemaining == 0 && newRemaining == 0)
+ return 0;
+
+ boolean different;
+ if (oldRemaining != newRemaining)
+ {
+ assert oldRemaining == 0 || newRemaining == 0; // one of them is null
+ different = true;
+ }
+ else
+ {
+ different = IntStream.range(0, oldRemaining).anyMatch(i -> oldValue.get(i) != newValue.get(i));
+ }
+
+ long bytesUsed = 0;
+ if (different)
+ {
+ var primaryKey = indexContext.hasClustering() ? indexContext.keyFactory().create(key, clustering)
+ : indexContext.keyFactory().create(key);
+ // update bounds because only rows with vectors are included in the key bounds,
+ // so if the vector was null before, we won't have included it
+ updateKeyBounds(primaryKey);
+
+ // make the changes in this order, so we don't have a window where the row is not in the index at all
+ if (newRemaining > 0)
+ bytesUsed += graph.add(newValue, primaryKey, OnHeapGraph.InvalidVectorBehavior.FAIL);
+ if (oldRemaining > 0)
+ bytesUsed -= graph.remove(oldValue, primaryKey);
+
+ // remove primary key if it's no longer indexed
+ if (newRemaining <= 0 && oldRemaining > 0)
+ primaryKeys.remove(primaryKey);
+ }
+ return bytesUsed;
+ }
+
+ private void updateKeyBounds(PrimaryKey primaryKey) {
+ if (minimumKey == null)
+ minimumKey = primaryKey;
+ else if (primaryKey.compareTo(minimumKey) < 0)
+ minimumKey = primaryKey;
+ if (maximumKey == null)
+ maximumKey = primaryKey;
+ else if (primaryKey.compareTo(maximumKey) > 0)
+ maximumKey = primaryKey;
+ }
+
+ @Override
+ public KeyRangeIterator search(QueryContext queryContext, Expression expr, AbstractBounds<PartitionPosition> keyRange)
+ {
+ assert expr.getOp() == Expression.IndexOperator.ANN : "Only ANN is supported for vector search, received " + expr.getOp();
+
+ VectorQueryContext vectorQueryContext = queryContext.vectorContext();
+
+ var buffer = expr.lower.value.raw;
+ float[] qv = TypeUtil.decomposeVector(indexContext, buffer);
+
+ Bits bits;
+ if (!RangeUtil.coversFullRing(keyRange))
+ {
+ // if left bound is MIN_BOUND or KEY_BOUND, we need to include all token-only PrimaryKeys with same token
+ boolean leftInclusive = keyRange.left.kind() != PartitionPosition.Kind.MAX_BOUND;
+ // if right bound is MAX_BOUND or KEY_BOUND, we need to include all token-only PrimaryKeys with same token
+ boolean rightInclusive = keyRange.right.kind() != PartitionPosition.Kind.MIN_BOUND;
+ // if right token is MAX (Long.MIN_VALUE), there is no upper bound
+ boolean isMaxToken = keyRange.right.getToken().isMinimum(); // max token
+
+ PrimaryKey left = indexContext.keyFactory().create(keyRange.left.getToken()); // lower bound
+ PrimaryKey right = isMaxToken ? null : indexContext.keyFactory().create(keyRange.right.getToken()); // upper bound
+
+ Set<PrimaryKey> resultKeys = isMaxToken ? primaryKeys.tailSet(left, leftInclusive) : primaryKeys.subSet(left, leftInclusive, right, rightInclusive);
+ if (!vectorQueryContext.getShadowedPrimaryKeys().isEmpty())
+ resultKeys = resultKeys.stream().filter(pk -> !vectorQueryContext.containsShadowedPrimaryKey(pk)).collect(Collectors.toSet());
+
+ if (resultKeys.isEmpty())
+ return KeyRangeIterator.empty();
+
+ int bruteForceRows = (int)(indexContext.getIndexWriterConfig().getMaximumNodeConnections() * Math.log(graph.size()));
+ if (resultKeys.size() < Math.max(vectorQueryContext.limit(), bruteForceRows))
+ return new ReorderingRangeIterator(new PriorityQueue<>(resultKeys));
+ else
+ bits = new KeyRangeFilteringBits(keyRange, vectorQueryContext.bitsetForShadowedPrimaryKeys(graph));
+ }
+ else
+ {
+ // partition/range deletion won't trigger index update, so we have to filter shadow primary keys in memtable index
+ bits = queryContext.vectorContext().bitsetForShadowedPrimaryKeys(graph);
+ }
+
+ var keyQueue = graph.search(qv, queryContext.vectorContext().limit(), bits);
+ if (keyQueue.isEmpty())
+ return KeyRangeIterator.empty();
+ return new ReorderingRangeIterator(keyQueue);
+ }
+
+ @Override
+ public KeyRangeIterator limitToTopResults(QueryContext context, KeyRangeIterator iterator, Expression exp)
+ {
+ Set<PrimaryKey> results = new HashSet<>();
+ while (iterator.hasNext())
+ {
+ var key = iterator.next();
+ if (!context.vectorContext().containsShadowedPrimaryKey(key))
+ results.add(key);
+ }
+
+ int maxBruteForceRows = Math.max(context.vectorContext().limit(), (int)(indexContext.getIndexWriterConfig().getMaximumNodeConnections() * Math.log(graph.size())));
+ if (results.size() <= maxBruteForceRows)
+ {
+ if (results.isEmpty())
+ return KeyRangeIterator.empty();
+ return new ReorderingRangeIterator(new PriorityQueue<>(results));
+ }
+
+ ByteBuffer buffer = exp.lower.value.raw;
+ float[] qv = TypeUtil.decomposeVector(indexContext, buffer);
+ var bits = new KeyFilteringBits(results);
+ var keyQueue = graph.search(qv, context.vectorContext().limit(), bits);
+ if (keyQueue.isEmpty())
+ return KeyRangeIterator.empty();
+ return new ReorderingRangeIterator(keyQueue);
+ }
+
+ @Override
+ public Iterator<Pair<ByteComparable, PrimaryKeys>> iterator()
+ {
+ // This method is only used when merging an in-memory index with a RowMapping. This is done a different
+ // way with the graph using the writeData method below.
+ throw new UnsupportedOperationException();
+ }
+
+ public SegmentMetadata.ComponentMetadataMap writeDirect(IndexDescriptor indexDescriptor,
+ IndexContext indexContext,
+ Function<PrimaryKey, Integer> postingTransformer) throws IOException
+ {
+ return graph.writeData(indexDescriptor, indexContext, postingTransformer);
+ }
+
+ @Override
+ public boolean isEmpty()
+ {
+ return graph.isEmpty();
+ }
+
+ @Nullable
+ @Override
+ public ByteBuffer getMinTerm()
+ {
+ return null;
+ }
+
+ @Nullable
+ @Override
+ public ByteBuffer getMaxTerm()
+ {
+ return null;
+ }
+
+ private class KeyRangeFilteringBits implements Bits
+ {
+ private final AbstractBounds<PartitionPosition> keyRange;
+ @Nullable
+ private final Bits bits;
+
+ public KeyRangeFilteringBits(AbstractBounds<PartitionPosition> keyRange, @Nullable Bits bits)
+ {
+ this.keyRange = keyRange;
+ this.bits = bits;
+ }
+
+ @Override
+ public boolean get(int ordinal)
+ {
+ if (bits != null && !bits.get(ordinal))
+ return false;
+
+ var keys = graph.keysFromOrdinal(ordinal);
+ return keys.stream().anyMatch(k -> keyRange.contains(k.partitionKey()));
+ }
+
+ @Override
+ public int length()
+ {
+ return graph.size();
+ }
+ }
+
+ private class ReorderingRangeIterator extends KeyRangeIterator
+ {
+ private final PriorityQueue<PrimaryKey> keyQueue;
+
+ ReorderingRangeIterator(PriorityQueue<PrimaryKey> keyQueue)
+ {
+ super(minimumKey, maximumKey, keyQueue.size());
+ this.keyQueue = keyQueue;
+ }
+
+ @Override
+ // VSTODO maybe we can abuse "current" to avoid having to pop and re-add the last skipped key
+ protected void performSkipTo(PrimaryKey nextKey)
Review Comment:
this is buggy, fix is in vsearch ebef8ea6bf56
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358100772
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -281,14 +288,29 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
Selectors selectors = selection.newSelectors(options);
AggregationSpecification aggregationSpec = getAggregationSpec(options);
- ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
- nowInSec, userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+ DataLimits limit = getDataLimits(userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+
+ // Handle additional validation for topK queries
+ if (restrictions.isTopK())
+ {
+ checkFalse(aggregationSpec != null, TOPK_AGGREGATION_ERROR);
+
+ checkFalse(options.getConsistency() != ConsistencyLevel.ONE &&
Review Comment:
As discussed on Slack conversation, we should do the same with the check for page size right below. So a page size below the query limit will produce a client warning and be ignored.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "michaelsembwever (via GitHub)" <gi...@apache.org>.
michaelsembwever commented on PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#issuecomment-1784080222
Committed with https://github.com/apache/cassandra/commit/949b760f5516c139591473038917247b1fd7f500
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357005728
##########
.build/parent-pom-template.xml:
##########
@@ -1182,6 +1182,11 @@
<artifactId>lucene-analysis-common</artifactId>
<version>9.7.0</version>
</dependency>
+ <dependency>
+ <groupId>io.github.jbellis</groupId>
+ <artifactId>jvector</artifactId>
+ <version>0.9.2</version>
Review Comment:
1.0.2 has a bunch of bug fixes, IIRC the api changed slightly from 0.9.x but it's an easy change
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "JeremiahDJordan (via GitHub)" <gi...@apache.org>.
JeremiahDJordan commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357060159
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -281,14 +288,29 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
Selectors selectors = selection.newSelectors(options);
AggregationSpecification aggregationSpec = getAggregationSpec(options);
- ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
- nowInSec, userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+ DataLimits limit = getDataLimits(userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+
+ // Handle additional validation for topK queries
+ if (restrictions.isTopK())
+ {
+ checkFalse(aggregationSpec != null, TOPK_AGGREGATION_ERROR);
+
+ checkFalse(options.getConsistency() != ConsistencyLevel.ONE &&
Review Comment:
Rather than failing queries not using LOCAL_ONE/ONE I think we should downgrade the CL to ONE/LOCAL_ONE (depending on if the incoming query used LOCAL_) and set a Client Warn that it happened.
It will be very painful for end users if they have to specially manage that ANN read queries use ONE when everything else in their app is using LOCAL_QUORUM...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357182902
##########
src/java/org/apache/cassandra/index/sai/plan/vector/VectorQueryController.java:
##########
@@ -0,0 +1,247 @@
+/*
+ * 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.index.sai.plan.vector;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.memtable.Memtable;
+import org.apache.cassandra.exceptions.QueryCancelledException;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.disk.SSTableRowIdKeyRangeIterator;
+import org.apache.cassandra.index.sai.disk.v1.postings.MergePostingList;
+import org.apache.cassandra.index.sai.disk.v1.vector.CheckpointingIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator;
+import org.apache.cassandra.index.sai.metrics.TableQueryMetrics;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.plan.StorageAttachedIndexQueryController;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.postings.PostingListIntersection;
+import org.apache.cassandra.index.sai.postings.PostingListUnion;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.utils.Throwables;
+
+public class VectorQueryController extends StorageAttachedIndexQueryController
+{
+ private static final Logger logger = LoggerFactory.getLogger(VectorQueryController.class);
+
+ public VectorQueryController(ColumnFamilyStore cfs,
+ ReadCommand command,
+ RowFilter filterOperation,
+ QueryContext queryContext,
+ TableQueryMetrics tableQueryMetrics)
+ {
+ super(cfs, command, filterOperation, queryContext, tableQueryMetrics);
+ }
+
+ /**
+ * Build a {@link KeyRangeIterator.Builder} from the given list of expressions by applying given operation (OR/AND).
+ * Building of such builder involves index search, results of which are persisted in the internal resources list
+ *
+ * @param expressions The expressions to build range iterator from (expressions with not results are ignored).
+ *
+ * @return range iterator builder based on given expressions and operation type.
+ */
+ public KeyRangeIterator.Builder getIndexQueryResults(Collection<Expression> expressions)
+ {
+ // FIXME having this at the expression level means that it only gets applied to Nodes at that level;
+ // moving it to ORDER BY will allow us to apply it correctly for other Node sub-trees
+ var annExpressionInHybridSearch = getAnnExpressionInHybridSearch(expressions);
+ boolean isAnnHybridSearch = annExpressionInHybridSearch != null;
+ if (isAnnHybridSearch)
+ expressions = expressions.stream().filter(e -> e != annExpressionInHybridSearch).collect(Collectors.toList());
+
+ // search memtable before referencing sstable indexes; otherwise we may miss newly flushed memtable index
+ var iteratorsByMemtable = expressions.stream()
+ .filter(expr -> !expr.context.isNotIndexed())
+ .flatMap(expr -> expr.context.getMemtableIndexManager().iteratorsForSearch(queryContext, expr, mergeRange).stream())
+ .collect(Collectors.groupingBy(pair -> pair.left, Collectors.mapping(pair -> pair.right, Collectors.toList())));
+
+ VectorQueryViewBuilder.QueryView queryView = new VectorQueryViewBuilder(expressions, mergeRange).build();
+ // in case of ANN query in hybrid search, we have to reference ANN sstable indexes separately because queryView doesn't include ANN sstable indexes
+ VectorQueryViewBuilder.QueryView annQueryViewInHybridSearch = isAnnHybridSearch ? new VectorQueryViewBuilder(Collections.singleton(annExpressionInHybridSearch), mergeRange).build() : null;
+
+ try
+ {
+ List<KeyRangeIterator> sstableIntersections = new ArrayList<>(queryView.view.size());
+ for (Map.Entry<SSTableReader, List<VectorQueryViewBuilder.IndexExpression>> viewEntry : queryView.view.entrySet())
+ {
+ PostingList postings = createRowIdIterator(viewEntry.getValue(), isAnnHybridSearch);
+ if (isAnnHybridSearch)
+ sstableIntersections.add(reorderAndLimitBySSTableRowIds(postings, viewEntry.getKey(), annQueryViewInHybridSearch));
+ else
+ {
+ PrimaryKeyMap.Factory pkFactory = viewEntry.getValue().iterator().next().index.getSSTableContext().primaryKeyMapFactory;
+ sstableIntersections.add(convertToPrimaryKeyIterator(pkFactory, postings));
+ }
+ }
+
+ List<KeyRangeIterator> memtableIntersections = iteratorsByMemtable.entrySet()
+ .stream()
+ .map(e -> {
+ // we need to do all the intersections at the index level, or ordering won't work
+ KeyRangeIterator it = buildIterator(e.getValue(), isAnnHybridSearch);
+ if (isAnnHybridSearch)
+ it = reorderAndLimitBy(it, e.getKey(), annExpressionInHybridSearch);
+ return it;
+ })
+ .collect(Collectors.toList());
+
+ Iterable<KeyRangeIterator> allIntersections = Iterables.concat(sstableIntersections, memtableIntersections);
Review Comment:
I don't think this is correct. Need to combine memtable results w/ sstable results at the expression level or you can miss matching rows that are split across multiple sources.
vsearch fix for this is 49a51690ca. it's a big diff though because you did a bunch of other refactoring here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356710478
##########
src/java/org/apache/cassandra/service/reads/range/ScanAllRangesCommandIterator.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.service.reads.range;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import com.google.common.base.Preconditions;
+
+import org.apache.cassandra.db.ConsistencyLevel;
+import org.apache.cassandra.db.EmptyIterators;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.PartitionRangeReadCommand;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.locator.EndpointsForRange;
+import org.apache.cassandra.locator.InetAddressAndPort;
+import org.apache.cassandra.locator.ReplicaPlan;
+import org.apache.cassandra.locator.ReplicaPlans;
+import org.apache.cassandra.net.Message;
+import org.apache.cassandra.net.MessagingService;
+import org.apache.cassandra.service.reads.DataResolver;
+import org.apache.cassandra.service.reads.ReadCallback;
+import org.apache.cassandra.service.reads.repair.NoopReadRepair;
+import org.apache.cassandra.tracing.Tracing;
+import org.apache.cassandra.utils.CloseableIterator;
+
+/**
+ * A custom {@link RangeCommandIterator} that queries all replicas required by consistency level at once with data range
+ * specify in {@link PartitionRangeReadCommand}.
+ *
+ * This is to speed up {@link Index.QueryPlan#isTopK()} queries that needs to find global top-k rows in the cluster, because
+ * existing {@link RangeCommandIterator} has to execute a top-k search per vnode range which is wasting resources.
+ */
+public class ScanAllRangesCommandIterator extends RangeCommandIterator
Review Comment:
We could also add a line saying that `ScanAllRangesCommandIterator` can only be used with CL=ONE/LOCAL_ONE. I understand we won't be able to use this optimization when we add support for higher CLs, making top-k queries quite expensive for vnodes. At least, not without further changes in `DataResolver`, RFP, RR, etc.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358203049
##########
src/java/org/apache/cassandra/db/marshal/ByteArrayAccessor.java:
##########
@@ -236,6 +236,19 @@ public float toFloat(byte[] value)
return ByteArrayUtil.getFloat(value, 0);
}
+ @Override
+ public float[] toFloatArray(byte[] value, int dimension)
Review Comment:
I agree. FWIW we are only using the `ByteBuffer` implementation atm and that does use `FloatBuffer.get`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1350437186
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1133,13 +1149,13 @@ public SelectStatement prepare(ClientState state, boolean forView) throws Invali
List<Selectable> selectables = RawSelector.toSelectables(selectClause, table);
boolean containsOnlyStaticColumns = selectOnlyStaticColumns(table, selectables);
- StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, containsOnlyStaticColumns, forView);
+ List<Ordering> orderings = getOrderings(table);
+ StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, orderings, containsOnlyStaticColumns, forView);
// If we order post-query, the sorted column needs to be in the ResultSet for sorting,
// even if we don't ultimately ship them to the client (CASSANDRA-4911).
- Map<ColumnMetadata, Boolean> orderingColumns = getOrderingColumns(table);
- Set<ColumnMetadata> resultSetOrderingColumns = restrictions.keyIsInRelation() ? orderingColumns.keySet()
- : Collections.emptySet();
+ Map<ColumnMetadata, Ordering> orderingColumns = getOrderingColumns(orderings);
+ Set<ColumnMetadata> resultSetOrderingColumns = getResultSetOrdering(restrictions, orderingColumns);
Review Comment:
Good catch, I have +1ed CASSANDRA-18892. I'll try to commit it tomorrow.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1364016948
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
Review Comment:
It's not so much of limitation, as a design decision. Splitting, the index up into N pieces slows down search by (almost) a factor of N.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1364114377
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
Review Comment:
That's fine to start with. Ideally we would flush tables with vector indexes to a single directory and leave others unaffected. Not sure how hard that is.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370603612
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -476,6 +507,38 @@ public boolean usesSecondaryIndexing()
return this.usesSecondaryIndexing;
}
+ /**
+ * This is a hack to push ordering down to indexes.
+ * Indexes are selected based on RowFilter only, so we need to turn orderings into restrictions
+ * so they end up in the row filter.
+ *
+ * @param orderings orderings from the select statement
+ * @return the {@link RestrictionSet} with the added orderings
+ */
+ private RestrictionSet addOrderingRestrictions(List<Ordering> orderings, RestrictionSet restrictionSet)
+ {
+ List<Ordering> annOrderings = orderings.stream().filter(o -> o.expression.hasNonClusteredOrdering()).collect(Collectors.toList());
+
+ if (annOrderings.size() > 1)
+ throw new InvalidRequestException("Cannot specify more than one ANN ordering");
+ else if (annOrderings.size() == 1)
+ {
+ if (orderings.size() > 1)
+ throw new InvalidRequestException("ANN ordering does not support secondary ordering");
Review Comment:
I think we are still missing this one in the last batch of commits.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370476474
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,35 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
Review Comment:
That documentation seems to contain a bunch of unreachable documents, if I'm understanding the links correctly. I'd say it has been accidentally committed. So probably it's better to ignore that by now and have a separate ticket for the doc.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1373043288
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,35 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
Review Comment:
I agree, let's deal with that separately. I have dropped a comment on CASSANDRA-18231.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1371609456
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
Review Comment:
I have added a validation to SAI to block creation of vector indexes on multiple data directories.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1336148245
##########
src/java/org/apache/cassandra/cql3/SingleColumnRelation.java:
##########
@@ -264,6 +265,16 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
return new SingleColumnRestriction.LikeRestriction(columnDef, operator, term);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ ColumnMetadata columnDef = table.getExistingColumn(entity);
+ if (!(columnDef.type instanceof VectorType))
Review Comment:
This code is redundant because we don't support ANN as a relation, only as an ordering. As such, I've removed this and all associated code.
##########
src/java/org/apache/cassandra/cql3/SingleColumnRelation.java:
##########
@@ -264,6 +265,16 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
return new SingleColumnRestriction.LikeRestriction(columnDef, operator, term);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ ColumnMetadata columnDef = table.getExistingColumn(entity);
+ if (!(columnDef.type instanceof VectorType))
+ throw invalidRequest("ANN is only supported against DENSE FLOAT32 columns");
Review Comment:
As previous message
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341331185
##########
src/java/org/apache/cassandra/cql3/Ordering.java:
##########
@@ -0,0 +1,187 @@
+/*
+ * 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;
+
+import org.apache.cassandra.cql3.restrictions.SingleColumnRestriction;
+import org.apache.cassandra.cql3.restrictions.SingleRestriction;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.TableMetadata;
+
+/**
+ * A single element of an ORDER BY clause.
+ * <code>ORDER BY ordering1 [, ordering2 [, ...]] </code>
+ * <p>
+ * An ordering comprises an expression that produces the values to compare against each other
+ * and a sorting direction (ASC, DESC).
+ */
+public class Ordering
+{
+ public final Expression expression;
+ public final Direction direction;
+
+ public Ordering(Expression expression, Direction direction)
+ {
+ this.expression = expression;
+ this.direction = direction;
+ }
+
+ public interface Expression
+ {
+ default boolean hasNonClusteredOrdering()
+ {
+ return false;
+ }
+
+ default SingleRestriction toRestriction()
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ ColumnMetadata getColumn();
+ }
+
+ /**
+ * Represents a single column in
+ * <code>ORDER BY column</code>
+ */
+ public static class SingleColumn implements Expression
+ {
+ public final ColumnMetadata column;
+
+ public SingleColumn(ColumnMetadata column)
+ {
+ this.column = column;
+ }
+
+ @Override
+ public ColumnMetadata getColumn()
+ {
+ return column;
+ }
+ }
+
+ /**
+ * An expression used in Approximate Nearest Neighbor ordering.
+ * <code>ORDER BY column ANN OF value</code>
+ */
+ public static class Ann implements Expression
+ {
+ final ColumnMetadata column;
+ final Term vectorValue;
+
+ public Ann(ColumnMetadata column, Term vectorValue)
+ {
+ this.column = column;
+ this.vectorValue = vectorValue;
+ }
+
+ @Override
+ public boolean hasNonClusteredOrdering()
+ {
+ return true;
+ }
+
+ @Override
+ public SingleRestriction toRestriction()
+ {
+ return new SingleColumnRestriction.AnnRestriction(column, vectorValue);
+ }
+
+ @Override
+ public ColumnMetadata getColumn()
+ {
+ return column;
+ }
+ }
+
+ public enum Direction
+ {ASC, DESC}
+
+
+ /**
+ * Represents the AST of a single element in the ORDER BY clause.
Review Comment:
I see, thanks. I don't see the changes in the JavaDoc. Maybe it could be changed to something like:
```suggestion
* Represents the ANTLR's abstract syntax tree of a single element in the {@code ORDER BY} clause.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1333481560
##########
src/java/org/apache/cassandra/index/sai/disk/format/IndexComponent.java:
##########
@@ -50,6 +50,11 @@ public enum IndexComponent
*/
TERMS_DATA("TermsData"),
+ /**
+ * Product Quantization store used to store compressed vectors for the vector index
+ */
+ PQ("PQ"),
Review Comment:
Maybe `PQ` is not the most friendly nor easy to investigate name when looking at the sstable files without context. We could extend it so it looks like the other components:
```suggestion
PRODUCT_QUANTIZATION("ProductQuantization"),
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1348796706
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/CassandraDiskAnn.java:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.PrimitiveIterator;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import io.github.jbellis.jvector.disk.CachingGraphIndex;
+import io.github.jbellis.jvector.disk.CompressedVectors;
+import io.github.jbellis.jvector.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.disk.ReaderSupplier;
+import io.github.jbellis.jvector.graph.GraphIndex;
+import io.github.jbellis.jvector.graph.GraphSearcher;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.NeighborSimilarity;
+import io.github.jbellis.jvector.graph.SearchResult.NodeScore;
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.disk.format.IndexComponent;
+import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
+import org.apache.cassandra.index.sai.disk.v1.postings.ReorderingPostingList;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.io.util.FileHandle;
+import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.io.util.RandomAccessReader;
+
+public class CassandraDiskAnn implements AutoCloseable
Review Comment:
Starting the class name with "Cassandra" doesn't give a lot of information, I think. Maybe this could be named `DiskANN`, `OnDiskANN`, or something like that?
Same for `CassandraOnHeapGraph`.
##########
src/java/org/apache/cassandra/index/sai/VectorQueryContext.java:
##########
@@ -0,0 +1,186 @@
+/*
+ * 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.index.sai;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+
+import io.github.jbellis.jvector.util.Bits;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.CassandraDiskAnn;
+import org.apache.cassandra.index.sai.disk.v1.vector.CassandraOnHeapGraph;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+
+
+/**
+ * This represents the state of a vector query. It is repsonsible for maintaining a list of any {@link PrimaryKey}s
+ * that have been updated or deleted during a search of the indexes.
+ * <p>
+ * The number of {@link #shadowedPrimaryKeys} is compared before and after a search is performed. If it changes, it
+ * means that a {@link PrimaryKey} was found to have been changed. In this case the whole search is repeated until the
+ * counts match.
+ * <p>
+ * When this process has completed, a {@link Bits} array is generated. This is used by the vector graph search to
+ * identify which nodes in the graph to include in the results.
+ */
+public class VectorQueryContext
+{
+ private TreeSet<PrimaryKey> shadowedPrimaryKeys; // allocate when needed
+ private int limit;
Review Comment:
Nit: can be `final`
##########
src/java/org/apache/cassandra/utils/ReadWriteLockedList.java:
##########
@@ -0,0 +1,115 @@
+/*
+ * 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.utils;
+
+import java.util.AbstractList;
+import java.util.List;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+public class ReadWriteLockedList<T> extends AbstractList<T>
Review Comment:
This entire class seems unused
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorSegmentationTest.java:
##########
@@ -0,0 +1,142 @@
+/*
+ * 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.index.sai.cql;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.Test;
+
+import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.db.marshal.FloatType;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentBuilder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class VectorSegmentationTest extends VectorTester
+{
+ private static final int dimension = 100;
+
+ @Test
+ public void testMultipleSegmentsForCreatingIndex() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+
+ int vectorCount = 100;
+ List<float[]> vectors = new ArrayList<>();
+ for (int row = 0; row < vectorCount; row++)
+ {
+ float[] vector = nextVector();
+ vectors.add(vector);
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", row, vector(vector));
+ }
+
+ flush();
+
+ SegmentBuilder.updateLastValidSegmentRowId(17); // 17 rows per segment
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ int limit = 35;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ @Test
+ public void testMultipleSegmentsForCompaction() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ List<float[]> vectors = new ArrayList<>();
+ int rowsPerSSTable = 10;
+ int sstables = 5;
+ int pk = 0;
+ for (int i = 0; i < sstables; i++)
+ {
+ for (int row = 0; row < rowsPerSSTable; row++)
+ {
+ float[] vector = nextVector();
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", pk++, vector(vector));
+ vectors.add(vector);
+ }
+
+ flush();
+ }
+
+ int limit = 30;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+
+
+ SegmentBuilder.updateLastValidSegmentRowId(11); // 11 rows per segment
+ compact();
+
+ queryVector = nextVector();
+ resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ resultVectors = getVectorsFromResult(resultSet);
+ recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ protected Vector<Float> vector(float[] values)
+ {
+ Float[] floats = new Float[values.length];
+ for (int i = 0; i < values.length; i++)
+ floats[i] = values[i];
+
+ return new Vector<>(floats);
+ }
+
+ private float[] nextVector()
Review Comment:
Nit: can be `static`
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java:
##########
@@ -0,0 +1,673 @@
+/*
+ * 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.index.sai.cql;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.junit.Test;
+
+import org.apache.cassandra.config.CassandraRelevantProperties;
+import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.service.ClientWarn;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class VectorTypeTest extends VectorTester
+{
+ private static final IPartitioner partitioner = Murmur3Partitioner.instance;
+
+ @Test
+ public void endToEndTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', [5.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (5, 'F', [6.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (6, 'G', [7.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (7, 'H', [8.0, 5.0, 6.0])");
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5");
+ assertThat(result).hasSize(5);
+
+ // some data that only lives in memtable
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (8, 'I', [9.0, 5.0, 6.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (9, 'J', [10.0, 6.0, 7.0])");
+ result = execute("SELECT * FROM %s ORDER BY val ann of [9.5, 5.5, 6.5] LIMIT 5");
+ assertContainsInt(result, "pk", 8);
+ assertContainsInt(result, "pk", 9);
+
+ // data from sstables
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertContainsInt(result, "pk", 1);
+ assertContainsInt(result, "pk", 2);
+ }
+
+ @Test
+ public void warningIsIssuedOnIndexCreation()
+ {
+ ClientWarn.instance.captureWarnings();
+ createTable("CREATE TABLE %s (pk int, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+ List<String> warnings = ClientWarn.instance.getWarnings();
+
+ assertTrue(warnings.size() > 0);
+ assertEquals(StorageAttachedIndex.VECTOR_USAGE_WARNING, warnings.get(0));
+ }
+
+ @Test
+ public void createIndexAfterInsertTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ flush();
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+ }
+
+ public static void assertContainsInt(UntypedResultSet result, String columnName, int columnValue)
+ {
+ for (UntypedResultSet.Row row : result)
+ {
+ if (row.has(columnName))
+ {
+ int value = row.getInt(columnName);
+ if (value == columnValue)
+ {
+ return;
+ }
+ }
+ }
+ throw new AssertionError("Result set does not contain a row with " + columnName + " = " + columnValue);
+ }
+
+ @Test
+ public void testTwoPredicates()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, b, v) VALUES (0, true, [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (1, true, [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (2, false, [3.0, 4.0, 5.0])");
+
+ // the vector given is closest to row 2, but we exclude that row because b=false
+ var result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ // VSTODO assert specific row keys
+ assertThat(result).hasSize(2);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testTwoPredicatesManyRows()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ for (int i = 0; i < 100; i++)
+ execute("INSERT INTO %s (pk, b, v) VALUES (?, true, ?)",
+ i, vector((float) i, (float) (i + 1), (float) (i + 2)));
+
+ var result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testThreePredicates()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, str text, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(str) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (0, true, [1.0, 2.0, 3.0], 'A')");
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (1, true, [2.0, 3.0, 4.0], 'B')");
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (2, false, [3.0, 4.0, 5.0], 'C')");
+
+ // the vector given is closest to row 2, but we exclude that row because b=false and str!='B'
+ var result = execute("SELECT * FROM %s WHERE b=true AND str='B' ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ // VSTODO assert specific row keys
+ assertThat(result).hasSize(1);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true AND str='B' ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void testSameVectorMultipleRows()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'A', [1.0, 2.0, 3.0])");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+ }
+
+ @Test
+ public void testQueryEmptyTable()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertThat(result).hasSize(0);
+ }
+
+ @Test
+ public void testQueryTableWithNulls()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', null)");
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertThat(result).hasSize(0);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [4.0, 5.0, 6.0])");
+ result = execute("SELECT pk FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertRows(result, row(1));
+ }
+
+ @Test
+ public void testLimitLessThanInsertedRowCount()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ // Insert more rows than the query limit
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [4.0, 5.0, 6.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [7.0, 8.0, 9.0])");
+
+ // Query with limit less than inserted row count
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testQueryMoreRowsThanInserted()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void changingOptionsTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ if (CassandraRelevantProperties.SAI_VECTOR_ALLOW_CUSTOM_PARAMETERS.getBoolean())
+ {
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = " +
+ "{'maximum_node_connections' : 10, 'construction_beam_width' : 200, 'similarity_function' : 'euclidean' }");
+ }
+ else
+ {
+ assertThatThrownBy(() -> createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = " +
+ "{'maximum_node_connections' : 10, 'construction_beam_width' : 200, 'similarity_function' : 'euclidean' }"))
+ .isInstanceOf(InvalidRequestException.class);
+ return;
+ }
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', [5.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (5, 'F', [6.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (6, 'G', [7.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (7, 'H', [8.0, 5.0, 6.0])");
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5");
+ assertThat(result).hasSize(5);
+ }
+
+ @Test
+ public void bindVariablesTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', ?)", vector(2.0f ,3.0f, 4.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', ?)", vector(4.0f, 5.0f, 6.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of ? LIMIT 3", vector(2.5f, 3.5f, 4.5f));
+ assertThat(result).hasSize(3);
+ }
+
+ @Test
+ public void intersectedSearcherTest()
+ {
+ // check that we correctly get back the two rows with str_val=B even when those are not
+ // the closest rows to the query vector
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', ?)", vector(2.0f ,3.0f, 4.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'B', ?)", vector(4.0f, 5.0f, 6.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', ?)", vector(5.0f, 6.0f, 7.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+
+ flush();
+ result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void nullVectorTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val) VALUES (1, 'B')"); // no vector
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val) VALUES (3, 'D')"); // no vector
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', ?)", vector(5.0f, 6.0f, 7.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(0);
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'A' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+
+ flush();
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(0);
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'A' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void lwtTest()
+ {
+ createTable("CREATE TABLE %s (p int, c int, v text, vec vector<float, 2>, PRIMARY KEY(p, c))");
+ createIndex("CREATE CUSTOM INDEX ON %s(vec) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (p, c, v) VALUES (?, ?, ?)", 0, 0, "test");
+ execute("INSERT INTO %s (p, c, v) VALUES (?, ?, ?)", 0, 1, "00112233445566");
+
+ execute("UPDATE %s SET v='00112233', vec=[0.9, 0.7] WHERE p = 0 AND c = 0 IF v = 'test'");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY vec ANN OF [0.1, 0.9] LIMIT 100");
+
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void twoVectorFieldsTest()
+ {
+ createTable("CREATE TABLE %s (pk int, v2 vector<float, 2>, v3 vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(v2) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v3) USING 'StorageAttachedIndex'");
+ }
+
+ @Test
+ public void primaryKeySearchTest()
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, 3>, i int, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ var N = 5;
+ for (int i = 0; i < N; i++)
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", i, vector(1.0f + i, 2.0f + i, 3.0f + i));
+
+ for (int i = 0; i < N; i++)
+ {
+ UntypedResultSet result = execute("SELECT pk FROM %s WHERE pk = ? ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2", i);
+ assertThat(result).hasSize(1);
+ assertRows(result, row(i));
+ }
+
+ flush();
+ for (int i = 0; i < N; i++)
+ {
+ UntypedResultSet result = execute("SELECT pk FROM %s WHERE pk = ? ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2", i);
+ assertThat(result).hasSize(1);
+ assertRows(result, row(i));
+ }
+ }
+
+ @Test
+ public void partitionKeySearchTest()
+ {
+ createTable("CREATE TABLE %s (partition int, row int, val vector<float, 2>, PRIMARY KEY(partition, row))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
+
+ var nPartitions = 5;
+ var rowsPerPartition = 10;
+ Map<Integer, List<float[]>> vectorsByPartition = new HashMap<>();
+
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ for (int j = 1; j <= rowsPerPartition; j++)
+ {
+ logger.debug("Inserting partition {} row {}: [{}, {}]", i, j, i, j);
+ execute("INSERT INTO %s (partition, row, val) VALUES (?, ?, ?)", i, j, vector((float) i, (float) j));
+ float[] vector = {(float) i, (float) j};
+ vectorsByPartition.computeIfAbsent(i, k -> new ArrayList<>()).add(vector);
+ }
+ }
+
+ var queryVector = vector(new float[] { 1.5f, 1.5f });
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition, row FROM %s WHERE partition = ? ORDER BY val ann of ? LIMIT 2", i, queryVector);
+ assertThat(result).hasSize(2);
+ assertRowsIgnoringOrder(result,
+ row(i, 1),
+ row(i, 2));
+ }
+
+ flush();
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition, row FROM %s WHERE partition = ? ORDER BY val ann of ? LIMIT 2", i, queryVector);
+ assertThat(result).hasSize(2);
+ assertRowsIgnoringOrder(result,
+ row(i, 1),
+ row(i, 2));
+ }
+ }
+
+ @Test
+ public void clusteringKeyIndexTest()
+ {
+ createTable("CREATE TABLE %s (pk int, ck vector<float, 2>, PRIMARY KEY(pk, ck))");
+ createIndex("CREATE CUSTOM INDEX ON %s(ck) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, ck) VALUES (1, [1.0, 2.0])");
+
+ assertRows(execute("SELECT * FROM %s ORDER BY ck ANN OF [1.0, 2.0] LIMIT 1"), row(1, vector(1.0F, 2.0F)));
+ }
+
+ @Test
+ public void rangeSearchTest() throws Throwable
+ {
+ createTable("CREATE TABLE %s (partition int, val vector<float, 2>, PRIMARY KEY(partition))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
+
+ var nPartitions = 100;
+ Map<Integer, float[]> vectorsByKey = new HashMap<>();
+
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ float[] vector = {(float) i, (float) i};
+ execute("INSERT INTO %s (partition, val) VALUES (?, ?)", i, vector(vector));
+ vectorsByKey.put(i, vector);
+ }
+
+ var queryVector = vector(new float[] { 1.5f, 1.5f });
+ CheckedFunction tester = () -> {
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithLowerBound(vectorsByKey.keySet(), i, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithLowerBound(vectorsByKey.keySet(), i, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithUpperBound(vectorsByKey.keySet(), i, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithUpperBound(vectorsByKey.keySet(), i, true));
+
+ for (int j = 1; j <= nPartitions; j++)
+ {
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) AND token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, true, j, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) AND token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, false, j, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) AND token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, true, j, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) AND token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, false, j, false));
+ }
+ }
+ };
+
+ tester.apply();
+
+ flush();
+
+ tester.apply();
+ }
+
+ private Collection<Integer> keys(UntypedResultSet result)
+ {
+ List<Integer> keys = new ArrayList<>(result.size());
+ for (UntypedResultSet.Row row : result)
+ keys.add(row.getInt("partition"));
+ return keys;
+ }
+
+ private Collection<Integer> keysWithLowerBound(Collection<Integer> keys, int leftKey, boolean leftInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getToken(Int32Type.instance.decompose(leftKey)), leftInclusive,
+ partitioner.getMaximumToken().getToken(), true);
+ }
+
+ private Collection<Integer> keysWithUpperBound(Collection<Integer> keys, int rightKey, boolean rightInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getMinimumToken().getToken(), true,
+ partitioner.getToken(Int32Type.instance.decompose(rightKey)), rightInclusive);
+ }
+
+ private Collection<Integer> keysInBounds(Collection<Integer> keys, int leftKey, boolean leftInclusive, int rightKey, boolean rightInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getToken(Int32Type.instance.decompose(leftKey)), leftInclusive,
+ partitioner.getToken(Int32Type.instance.decompose(rightKey)), rightInclusive);
+ }
+
+ private Collection<Integer> keysInTokenRange(Collection<Integer> keys, Token leftToken, boolean leftInclusive, Token rightToken, boolean rightInclusive)
+ {
+ long left = leftToken.getLongValue();
+ long right = rightToken.getLongValue();
+ return keys.stream()
+ .filter(k -> {
+ long t = partitioner.getToken(Int32Type.instance.decompose(k)).getLongValue();
+ return (left < t || left == t && leftInclusive) && (t < right || t == right && rightInclusive);
+ }).collect(Collectors.toSet());
+ }
+
+ @Test
+ public void selectFloatVectorFunctions()
+ {
+ createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, value vector<float, 2>)");
+
+ // basic functionality
+ Vector<Float> q = vector(1f, 2f);
+ execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(1f, 2f));
+ execute("SELECT similarity_cosine(value, value) FROM %s WHERE pk=0");
+
+ // type inference checks
+ var result = execute("SELECT similarity_cosine(value, ?) FROM %s WHERe pk=0", q);
+ assertRows(result, row(1f));
+ result = execute("SELECT similarity_euclidean(value, ?) FROM %s WHERe pk=0", q);
+ assertRows(result, row(1f));
+ execute("SELECT similarity_cosine(?, value) FROM %s WHERE pk=0", q);
+ assertThatThrownBy(() -> execute("SELECT similarity_cosine(?, ?) FROM %s WHERE pk=0", q, q))
+ .hasMessageContaining("Cannot infer type of argument ?");
+
+ // with explicit typing
+ execute("SELECT similarity_cosine((vector<float, 2>) ?, ?) FROM %s WHERE pk=0", q, q);
+ execute("SELECT similarity_cosine(?, (vector<float, 2>) ?) FROM %s WHERE pk=0", q, q);
+ execute("SELECT similarity_cosine((vector<float, 2>) ?, (vector<float, 2>) ?) FROM %s WHERE pk=0", q, q);
+ }
+
+ @Test
+ public void selectSimilarityWithAnn()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ Vector<Float> q = vector(1.5f, 2.5f, 3.5f);
+ var result = execute("SELECT str_val, similarity_cosine(val, ?) FROM %s ORDER BY val ANN OF ? LIMIT 2",
+ q, q);
+
+ assertRowsIgnoringOrder(result,
+ row("A", 0.9987074f),
+ row("B", 0.9993764f));
+ }
+
+ @Test
+ public void castedTerminalFloatVectorFunctions()
+ {
+ createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, value vector<float, 2>)");
+
+ execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(1f, 2f));
+ execute("SELECT similarity_cosine(value, (vector<float, 2>) [1.0, 1.0]) FROM %s WHERE pk=0");
+ execute("SELECT similarity_cosine((vector<float, 2>) [1.0, 1.0], value) FROM %s WHERE pk=0");
+ execute("SELECT similarity_cosine((vector<float, 2>) [1.0, 1.0], (vector<float, 2>) [1.0, 1.0]) FROM %s WHERE pk=0");
+ }
+
+ @Test
+ public void inferredTerminalFloatVectorFunctions() throws Throwable
Review Comment:
It seems there is some overlapping between these two tests and `VectorFctsTest`
##########
test/unit/org/apache/cassandra/inject/injections.md:
##########
@@ -0,0 +1,356 @@
+<!---
+ 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.
+-->
+
+## Injecting hooks dynamically in testing
+
+The testing infrastructure is equipped with Byteman, which allows to
+inject hooks into the running code. In other words, at any point in
+the test, you can tell the particular nodes, to do something at exact
+location in the code - like you would set a conditional breakpoint
+and when the debugger breaks, you can check something or perform some
+additional actions. In short, it is particularly useful to:
+
+- synchronizing the test code and the nodes
+- synchronizing the nodes between each other
+- running actions step by step
+- counting invocations
+- tracing invocations
+- force throwing exceptions
+- and combinations of the above...
+
+Here you can find a short introduction and some examples how can it be
+used.
+
+### Basics
+
+Byteman works on rules which are the units defining what, when, and
+where should be invoked. That is, a single rule includes: the invoke
+point, bindings, run condition and actions.
+
+**Invoke point** is a location in the code where the actions should
+be hooked. Usually it is a class and method, but it can be specified
+more precisely, like - entering method, exiting method, invoking some
+method inside, particular line of code and so on.
+
+**Bindings** are just some constant definitions which can be used then
+in run condition and actions.
+
+**Run condition** is a logical expression which determines whether
+the rule should be invoked.
+
+**Actions** are the statements to be invoked.
+
+In DSE, we have `Rule` class which reflects a single Byteman rule.
Review Comment:
This file contains multiple references to DSE.
##########
src/java/org/apache/cassandra/io/util/MmappedRegions.java:
##########
@@ -268,6 +270,18 @@ public ByteBuffer buffer()
return buffer.duplicate();
}
+ public FloatBuffer floatBuffer()
+ {
+ // this does an implicit duplicate(), so we need to expose it directly to avoid doing it twice unnecessarily
+ return buffer.asFloatBuffer();
+ }
+
+ public IntBuffer intBuffer()
Review Comment:
Nit: Add `@Override`
##########
src/java/org/apache/cassandra/index/sai/plan/Expression.java:
##########
@@ -172,6 +178,9 @@ public Expression add(Operator op, ByteBuffer value)
*/
public boolean isSatisfiedBy(ByteBuffer columnValue)
{
+ if (validator.isVector())
Review Comment:
Nit: A brief comment on why we do this could be useful for readers without context.
##########
src/java/org/apache/cassandra/index/sai/VectorQueryContext.java:
##########
@@ -0,0 +1,186 @@
+/*
+ * 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.index.sai;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+
+import io.github.jbellis.jvector.util.Bits;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.CassandraDiskAnn;
+import org.apache.cassandra.index.sai.disk.v1.vector.CassandraOnHeapGraph;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+
+
+/**
+ * This represents the state of a vector query. It is repsonsible for maintaining a list of any {@link PrimaryKey}s
+ * that have been updated or deleted during a search of the indexes.
+ * <p>
+ * The number of {@link #shadowedPrimaryKeys} is compared before and after a search is performed. If it changes, it
+ * means that a {@link PrimaryKey} was found to have been changed. In this case the whole search is repeated until the
+ * counts match.
+ * <p>
+ * When this process has completed, a {@link Bits} array is generated. This is used by the vector graph search to
+ * identify which nodes in the graph to include in the results.
+ */
+public class VectorQueryContext
+{
+ private TreeSet<PrimaryKey> shadowedPrimaryKeys; // allocate when needed
Review Comment:
It might me helpful to add a brief comment saying what's a shadowed primary key.
##########
src/java/org/apache/cassandra/io/util/MmappedRegions.java:
##########
@@ -268,6 +270,18 @@ public ByteBuffer buffer()
return buffer.duplicate();
}
+ public FloatBuffer floatBuffer()
Review Comment:
Nit: Add `@Override`
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/OnDiskOrdinalsMap.java:
##########
@@ -0,0 +1,170 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
+
+import com.google.common.base.Preconditions;
+
+import io.github.jbellis.jvector.util.Bits;
+import org.apache.cassandra.io.util.FileHandle;
+import org.apache.cassandra.io.util.RandomAccessReader;
+
+public class OnDiskOrdinalsMap
+{
+ private final FileHandle fh;
+ private final long ordToRowOffset;
+ private final long segmentEnd;
+ private final int size;
+ // the offset where we switch from recording ordinal -> rows, to row -> ordinal
+ private final long rowOrdinalOffset;
+ private final Set<Integer> deletedOrdinals;
+
+ public OnDiskOrdinalsMap(FileHandle fh, long segmentOffset, long segmentLength)
+ {
+ deletedOrdinals = new HashSet<>();
+
+ this.segmentEnd = segmentOffset + segmentLength;
+ this.fh = fh;
+ try (var reader = fh.createReader())
+ {
+ reader.seek(segmentOffset);
+ int deletedCount = reader.readInt();
+ for (var i = 0; i < deletedCount; i++)
+ {
+ deletedOrdinals.add(reader.readInt());
+ }
+
+ this.ordToRowOffset = reader.getFilePointer();
+ this.size = reader.readInt();
+ reader.seek(segmentEnd - 8);
+ this.rowOrdinalOffset = reader.readLong();
+ assert rowOrdinalOffset < segmentEnd : "rowOrdinalOffset " + rowOrdinalOffset + " is not less than segmentEnd " + segmentEnd;
+ }
+ catch (Exception e)
+ {
+ throw new RuntimeException("Error initializing OnDiskOrdinalsMap at segment " + segmentOffset, e);
+ }
+ }
+
+ public RowIdsView getRowIdsView()
+ {
+ return new RowIdsView();
+ }
+
+ public Bits ignoringDeleted(Bits acceptBits)
+ {
+ return BitsUtil.bitsIgnoringDeleted(acceptBits, deletedOrdinals);
+ }
+
+ public class RowIdsView implements AutoCloseable
+ {
+ final RandomAccessReader reader = fh.createReader();
+
+ public int[] getSegmentRowIdsMatching(int vectorOrdinal) throws IOException
+ {
+ Preconditions.checkArgument(vectorOrdinal < size, "vectorOrdinal %s is out of bounds %s", vectorOrdinal, size);
+
+ // read index entry
+ try
+ {
+ reader.seek(ordToRowOffset + 4L + vectorOrdinal * 8L);
+ }
+ catch (Exception e)
+ {
+ throw new RuntimeException(String.format("Error seeking to index offset for ordinal %d with ordToRowOffset %d",
+ vectorOrdinal, ordToRowOffset), e);
+ }
+ var offset = reader.readLong();
+ // seek to and read rowIds
+ try
+ {
+ reader.seek(offset);
+ }
+ catch (Exception e)
+ {
+ throw new RuntimeException(String.format("Error seeking to rowIds offset for ordinal %d with ordToRowOffset %d",
+ vectorOrdinal, ordToRowOffset), e);
+ }
+ var postingsSize = reader.readInt();
+ var rowIds = new int[postingsSize];
+ for (var i = 0; i < rowIds.length; i++)
+ {
+ rowIds[i] = reader.readInt();
+ }
+ return rowIds;
+ }
+
+ @Override
+ public void close()
+ {
+ reader.close();
+ }
+ }
+
+ public OrdinalsView getOrdinalsView()
+ {
+ return new OrdinalsView();
+ }
+
+ public class OrdinalsView implements AutoCloseable
+ {
+ final RandomAccessReader reader = fh.createReader();
+ private final long high = (segmentEnd - 8 - rowOrdinalOffset) / 8;
+
+ /**
+ * @return order if given row id is found; otherwise return -1
+ */
+ public int getOrdinalForRowId(int rowId) throws IOException
+ {
+ // Compute the offset of the start of the rowId to vectorOrdinal mapping
+ long index = DiskBinarySearch.searchInt(0, Math.toIntExact(high), rowId, i -> {
+ try
+ {
+ long offset = rowOrdinalOffset + i * 8;
+ reader.seek(offset);
+ return reader.readInt();
+ }
+ catch (IOException e)
+ {
+ throw new RuntimeException(e);
+ }
+ });
+
+ // not found
+ if (index < 0)
+ return -1;
+
+ return reader.readInt();
+ }
+
+ @Override
+ public void close()
+ {
+ reader.close();
+ }
+ }
+
+ public void close()
Review Comment:
Nit: Maybe this class could implement `Closeable`, so this is annotated with override and more visible for callers?
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,476 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final String INVALID_LIMIT_MESSAGE = "Use of ANN OF in an ORDER BY clause requires a LIMIT that is not greater than 1000";
Review Comment:
Nit: unused constant
##########
src/java/org/apache/cassandra/index/sai/plan/StorageAttachedIndexSearcher.java:
##########
@@ -105,7 +104,25 @@ public PartitionIterator filterReplicaFilteringProtection(PartitionIterator full
@Override
public UnfilteredPartitionIterator search(ReadExecutionController executionController) throws RequestTimeoutException
{
- return new ResultRetriever(queryController, executionController, queryContext, keyFactory);
+ if (!command.isTopK())
+ return new ResultRetriever(queryController, executionController, queryContext, false);
+ else
+ {
+ Supplier<ResultRetriever> resultSupplier = () -> new ResultRetriever(queryController, executionController, queryContext, true);
+
+ // VSTODO performance: if there is shadowed primary keys, we have to at least query twice.
+ // First time to find out there are shawdow keys, second time to find out there are no more shadow keys.
Review Comment:
```suggestion
// First time to find out there are shadow keys, second time to find out there are no more shadow keys.
```
##########
src/java/org/apache/cassandra/index/sai/memory/MemtableOrdering.java:
##########
@@ -0,0 +1,43 @@
+/*
+ * 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.index.sai.memory;
+
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+
+/***
+ * Analogue of SegmentOrdering, but for memtables.
Review Comment:
```suggestion
* Analogue of {@link org.apache.cassandra.index.sai.disk.v1.segment.SegmentOrdering}, but for memtables.
```
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorSegmentationTest.java:
##########
@@ -0,0 +1,142 @@
+/*
+ * 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.index.sai.cql;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.Test;
+
+import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.db.marshal.FloatType;
+import org.apache.cassandra.db.marshal.VectorType;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentBuilder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class VectorSegmentationTest extends VectorTester
+{
+ private static final int dimension = 100;
+
+ @Test
+ public void testMultipleSegmentsForCreatingIndex() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+
+ int vectorCount = 100;
+ List<float[]> vectors = new ArrayList<>();
+ for (int row = 0; row < vectorCount; row++)
+ {
+ float[] vector = nextVector();
+ vectors.add(vector);
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", row, vector(vector));
+ }
+
+ flush();
+
+ SegmentBuilder.updateLastValidSegmentRowId(17); // 17 rows per segment
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ int limit = 35;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ @Test
+ public void testMultipleSegmentsForCompaction() throws Throwable
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, " + dimension + ">, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ List<float[]> vectors = new ArrayList<>();
+ int rowsPerSSTable = 10;
+ int sstables = 5;
+ int pk = 0;
+ for (int i = 0; i < sstables; i++)
+ {
+ for (int row = 0; row < rowsPerSSTable; row++)
+ {
+ float[] vector = nextVector();
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", pk++, vector(vector));
+ vectors.add(vector);
+ }
+
+ flush();
+ }
+
+ int limit = 30;
+ float[] queryVector = nextVector();
+ UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ List<float[]> resultVectors = getVectorsFromResult(resultSet);
+ double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+
+
+ SegmentBuilder.updateLastValidSegmentRowId(11); // 11 rows per segment
+ compact();
+
+ queryVector = nextVector();
+ resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
+ assertThat(resultSet.size()).isEqualTo(limit);
+
+ resultVectors = getVectorsFromResult(resultSet);
+ recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
+ assertThat(recall).isGreaterThanOrEqualTo(0.99);
+ }
+
+ protected Vector<Float> vector(float[] values)
+ {
+ Float[] floats = new Float[values.length];
+ for (int i = 0; i < values.length; i++)
+ floats[i] = values[i];
+
+ return new Vector<>(floats);
+ }
+
+ private float[] nextVector()
+ {
+ float[] rawVector = new float[dimension];
+ for (int i = 0; i < dimension; i++)
+ {
+ rawVector[i] = getRandom().nextFloat();
+ }
+ return rawVector;
+ }
+
+ private List<float[]> getVectorsFromResult(UntypedResultSet result)
Review Comment:
Nit: can be `static`
##########
test/unit/org/apache/cassandra/index/sai/utils/Glove.java:
##########
@@ -0,0 +1,112 @@
+/*
+ * 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.index.sai.utils;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class Glove
Review Comment:
This could have some class JavaDoc
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1360986023
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1080,19 +1139,26 @@ public void processPartition(RowIterator partition, QueryOptions options, Result
private boolean needsPostQueryOrdering()
{
- // We need post-query ordering only for queries with IN on the partition key and an ORDER BY.
- return restrictions.keyIsInRelation() && !parameters.orderings.isEmpty();
+ // We need post-query ordering only for queries with IN on the partition key and an ORDER BY or index restriction reordering
+ return restrictions.keyIsInRelation() && !parameters.orderings.isEmpty() || needIndexOrdering();
+ }
+
+ private boolean needIndexOrdering()
+ {
+ return orderingComparator != null && orderingComparator.indexOrdering();
}
/**
* Orders results when multiple keys are selected (using IN)
*/
- private void orderResults(ResultSet cqlRows)
+ private void orderResults(ResultSet cqlRows, QueryOptions options)
{
if (cqlRows.size() == 0 || !needsPostQueryOrdering())
return;
- Collections.sort(cqlRows.rows, orderingComparator);
+ Comparator<List<ByteBuffer>> comparator = orderingComparator.prepareFor(table, getRowFilter(options), options);
+ if (comparator != null)
+ Collections.sort(cqlRows.rows, comparator);
Review Comment:
Nit: this can be slightly simplified:
```suggestion
cqlRows.rows.sort(comparator);
```
Also, it might be worth adding a comment saying that this orders first by the indexed column, then by primary key. This guarantee will probably be useful if we support paging in the future.
The order is that way even when we only consider the indexed column when sorting here. That's because the replica results have been returned ordered by primary key. The stability of the used sorting algorithm (merge sort) ensures that rows with the same value-based index score preserve the original primary key order.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1368779258
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -274,10 +292,54 @@ else if (relation.isLIKE())
throw invalidRequest("Non PRIMARY KEY columns found in where clause: %s ",
Joiner.on(", ").join(nonPrimaryKeyColumns));
}
+
+ var annRestriction = Streams.stream(nonPrimaryKeyRestrictions).filter(SingleRestriction::isANN).findFirst();
+ if (annRestriction.isPresent())
+ {
+ // If there is an ANN restriction then it must be for a vector<float, n> column, and it must have an index
+ var annColumn = annRestriction.get().getFirstColumn();
+
+ if (!annColumn.type.isVector() || !(((VectorType<?>)annColumn.type).elementType instanceof FloatType))
+ throw invalidRequest(StatementRestrictions.ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE);
+ if (indexRegistry == null || indexRegistry.listIndexes().stream().noneMatch(i -> i.dependsOn(annColumn)))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEX_MESSAGE);
+ // We do not allow ANN queries using partition key restrictions that need filtering
+ if (partitionKeyRestrictions.needFiltering(table))
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE);
+ // We do not allow ANN query filtering using non-indexed columns
+ var nonAnnColumns = Streams.stream(nonPrimaryKeyRestrictions).filter(r -> !r.isANN()).map(r -> r.getFirstColumn()).collect(Collectors.toList());
+ var clusteringColumns = clusteringColumnsRestrictions.getColumnDefinitions();
+ if (!nonAnnColumns.isEmpty() || !clusteringColumns.isEmpty())
+ {
+ var nonIndexedColumns = Stream.concat(nonAnnColumns.stream(), clusteringColumns.stream())
+ .filter(c -> indexRegistry == null || indexRegistry.listIndexes()
+ .stream()
+ .noneMatch(i -> i.dependsOn(c)))
+ .collect(Collectors.toList());
+ if (!nonIndexedColumns.isEmpty())
+ throw invalidRequest(StatementRestrictions.ANN_REQUIRES_INDEXED_FILTERING_MESSAGE);
Review Comment:
Although it doesn't need `ALLOW FILTERING` and it doesn't get filtered by `applyIndexFilter`, it does get filtered earlier by `QueryController.doesNotSelect`. This would have the same issues as any other non-indexed filtering.
I feel that we should leave this restriction in place for the initial release and add it to the list of restrictions to fix later.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370280849
##########
src/java/org/apache/cassandra/index/sai/plan/VectorTopKProcessor.java:
##########
@@ -0,0 +1,198 @@
+/*
+ * 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.index.sai.plan;
+
+import java.nio.ByteBuffer;
+import java.util.Comparator;
+import java.util.Optional;
+import java.util.PriorityQueue;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import javax.annotation.Nullable;
+
+import com.google.common.base.Preconditions;
+import org.apache.commons.lang3.tuple.Triple;
+
+import org.apache.cassandra.cql3.Operator;
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.partitions.BasePartitionIterator;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.rows.BaseRowIterator;
+import org.apache.cassandra.db.rows.Row;
+import org.apache.cassandra.db.rows.Unfiltered;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.index.SecondaryIndexManager;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.index.sai.utils.InMemoryPartitionIterator;
+import org.apache.cassandra.index.sai.utils.InMemoryUnfilteredPartitionIterator;
+import org.apache.cassandra.index.sai.utils.PartitionInfo;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.utils.FBUtilities;
+import org.apache.cassandra.utils.Pair;
+
+/**
+ * Processor that scans all rows from given partitions and selects rows with top-k scores based on vector indexes.
+ * <p>
+ * This processor performs the following steps:
+ * - collect rows with score into PriorityQueue that sorts rows based on score. If there are multiple vector indexes,
Review Comment:
```suggestion
* - collect rows with score into a {@link PriorityQueue} that sorts rows based on score. If there are multiple vector indexes,
```
##########
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeListIterator.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.index.sai.iterators;
+
+import java.util.List;
+
+import com.google.common.collect.Iterators;
+import com.google.common.collect.PeekingIterator;
+
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+
+/**
+ * A {@link KeyRangeIterator} that iterates over a list of {@link PrimaryKey}s without modifying the underlying list.
+ */
+public class KeyRangeListIterator extends KeyRangeIterator
+{
+ private final PeekingIterator<PrimaryKey> keyQueue;
+
+ /**
+ * Create a new {@link KeyRangeListIterator} that iterates over the provided list of keys.
Review Comment:
```suggestion
* Create a new {@link KeyRangeListIterator} that iterates over the provided list of keys.
*
```
##########
src/java/org/apache/cassandra/index/sai/disk/v1/postings/VectorPostingList.java:
##########
@@ -0,0 +1,75 @@
+/*
+ * 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.index.sai.disk.v1.postings;
+
+import java.io.IOException;
+import java.util.PrimitiveIterator;
+
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.lucene.util.LongHeap;
+
+/**
+ * A posting list for ANN search results. Transforms result from similarity order to rowId order.
Review Comment:
```suggestion
* A {@link PostingList} for ANN search results. Transforms result from similarity order to row ID order.
```
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
Review Comment:
We still have to get rid of the mention of Converged Cassandra.
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,35 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
Review Comment:
I have just realized that the documentation for SAI committed by CASSANDRA-18231 a month ago contains documentation about vector search. We should verify that the doc matches this implementation, and probably add the list of query limitations to it, probably [here](https://github.com/apache/cassandra/blob/trunk/doc/modules/cassandra/partials/vector-search/query-vector-data.adoc).
##########
src/java/org/apache/cassandra/index/sai/plan/VectorTopKProcessor.java:
##########
@@ -0,0 +1,198 @@
+/*
+ * 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.index.sai.plan;
+
+import java.nio.ByteBuffer;
+import java.util.Comparator;
+import java.util.Optional;
+import java.util.PriorityQueue;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import javax.annotation.Nullable;
+
+import com.google.common.base.Preconditions;
+import org.apache.commons.lang3.tuple.Triple;
+
+import org.apache.cassandra.cql3.Operator;
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.partitions.BasePartitionIterator;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
+import org.apache.cassandra.db.rows.BaseRowIterator;
+import org.apache.cassandra.db.rows.Row;
+import org.apache.cassandra.db.rows.Unfiltered;
+import org.apache.cassandra.index.Index;
+import org.apache.cassandra.index.SecondaryIndexManager;
+import org.apache.cassandra.index.sai.IndexContext;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.index.sai.utils.InMemoryPartitionIterator;
+import org.apache.cassandra.index.sai.utils.InMemoryUnfilteredPartitionIterator;
+import org.apache.cassandra.index.sai.utils.PartitionInfo;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.utils.FBUtilities;
+import org.apache.cassandra.utils.Pair;
+
+/**
+ * Processor that scans all rows from given partitions and selects rows with top-k scores based on vector indexes.
+ * <p>
+ * This processor performs the following steps:
+ * - collect rows with score into PriorityQueue that sorts rows based on score. If there are multiple vector indexes,
+ * the final score is the sum of all vector index scores.
+ * - remove rows with the lowest scores from PQ if PQ size exceeds limit
+ * - return rows from PQ in primary key order to client
+ * <p>
+ * Note that recall will be lower with paging, because:
+ * - page size is used as limit
+ * - for the first query, coordinator returns global top page-size rows within entire ring
+ * - for the subsequent queries, coordinators returns global top page-size rows withom range from last-returned-row to max token
Review Comment:
We don't support paging, maybe this comes from a previous attempt to support it?
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+import java.util.List;
+
+import org.apache.cassandra.dht.AbstractBounds;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+
+/**
+ * A {@link SegmentOrdering} orders and limits a list of {@link PrimaryKey}s.
+ * <p>
+ * When using {@link SegmentOrdering} there are several steps to
+ * build the list of Primary Keys to be ordered and limited:
+ * <p>
+ * 1. Find all primary keys that match each non-ordering query predicate.
+ * 2. Union and intersect the results of step 1 to build a single {@link KeyRangeIterator}
+ * ordered by {@link PrimaryKey}.
+ * 3. Filter out any shadowed primary keys.
+ * 4. Fan the primary keys from step 3 out to each sstable segment to order and limit each
+ * list of primary keys.
+ * <p>
+ * SegmentOrdering handles the fourth step.
+ * <p>
+ * Note: a segment ordering is only used when a query has both ordering and non-ordering predicates.
+ * Where a query has only ordering predicates, the ordering is handled by the
Review Comment:
```suggestion
* Where a query has only ordering predicates, the ordering is handled by
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370489241
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +71,22 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row ID for a given token or a negative value if not found
+ *
+ * @param token the {@link Token} to lookup
+ * @return the first row ID associated with the {@link Token} or a negative value
+ */
+ long ceiling(Token token);
+
+ /**
+ * Returns the last row ID for a given partition range or the previous row ID if there was no match for the partition range
Review Comment:
Yes, I think the original code uses something `the row ID of the least PK greater than or equal to the given PK` for ceiling, and `the row ID of the greatest PK less than or equal to the given PK`, which sounds clearer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342759209
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -40,16 +41,21 @@
import java.util.stream.Collectors;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.Futures; // checkstyle: permit this import
import com.google.common.util.concurrent.ListenableFuture; // checkstyle: permit this import
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
Review Comment:
There is still a Lucene import on `VectorFctsTest`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1349576732
##########
src/java/org/apache/cassandra/db/marshal/VectorType.java:
##########
@@ -365,6 +366,7 @@ public abstract class VectorSerializer extends TypeSerializer<List<T>>
public abstract <V> List<V> split(V buffer, ValueAccessor<V> accessor);
public abstract <V> V serializeRaw(List<V> elements, ValueAccessor<V> accessor);
+ public abstract float[] deserializeFloatArray(ByteBuffer input);
Review Comment:
I see, but `deserializeFloatArray ` is an instance method with access to dimension. It produces the same output for the same input as `composeAsFloat` unless the input is of a different dimension than the type instance it belongs to, I think. Can this ever happen, that we call it with an input of different size? If that's the case, it should probably be a static method. If not, then one of the implementations seems redundant.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1350176074
##########
src/java/org/apache/cassandra/index/sai/IndexContext.java:
##########
@@ -354,6 +371,11 @@ public boolean isLiteral()
return TypeUtil.isLiteral(getValidator());
}
+ public boolean isVector()
+ {
+ return getValidator().isVector() && ((VectorType<?>)getValidator()).elementType instanceof FloatType;
Review Comment:
Yes, that makes sense. I guess that if we were to ever allow indexing vectors as literal one might consider making it the default, so the most basic equality indexing is the default. But ANN seems much more useful in practice, so I guess it's better to have it as the default.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357056488
##########
src/java/org/apache/cassandra/db/filter/RowFilter.java:
##########
@@ -698,6 +698,7 @@ public boolean isSatisfiedBy(TableMetadata metadata, DecoratedKey partitionKey,
case LIKE_SUFFIX:
case LIKE_CONTAINS:
case LIKE_MATCHES:
+ case ANN:
{
assert !column.isComplex() : "Only CONTAINS and CONTAINS_KEY are supported for 'complex' types";
Review Comment:
not part of this patch but this is a poor error message, "complex" is an internal thing, can we slip in a quick improvement?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358195155
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -63,6 +69,14 @@ public final class StatementRestrictions
"Executing this query despite the performance unpredictability with ALLOW FILTERING has been disabled " +
"by the allow_filtering_enabled property in cassandra.yaml";
+ public static final String ANN_REQUIRES_INDEX_MESSAGE = "ANN ordering by vector requires the column to be indexed";
+
+ public static final String VECTOR_INDEXES_ANN_ONLY_MESSAGE = "Vector indexes only support ANN queries";
+
+ public static final String ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE = "ANN ordering is only supported on float vector indexes";
+
+ public static final String ANN_REQUIRES_INDEXED_FILTERING_MESSAGE = "ANN ordered queries can only be filtered with indexed columns";
Review Comment:
Unfortunately, the `StatementRestrictions` implementation is very different on ds vsearch but I have picked up the improved error message.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358199942
##########
src/java/org/apache/cassandra/db/filter/RowFilter.java:
##########
@@ -698,6 +698,7 @@ public boolean isSatisfiedBy(TableMetadata metadata, DecoratedKey partitionKey,
case LIKE_SUFFIX:
case LIKE_CONTAINS:
case LIKE_MATCHES:
+ case ANN:
{
assert !column.isComplex() : "Only CONTAINS and CONTAINS_KEY are supported for 'complex' types";
Review Comment:
Agreed. I've changed the `'complex'` to `collection`. Seems the simplest change and makes it much clearer to the user.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358483415
##########
src/java/org/apache/cassandra/index/sai/plan/vector/VectorQueryController.java:
##########
@@ -0,0 +1,247 @@
+/*
+ * 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.index.sai.plan.vector;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.memtable.Memtable;
+import org.apache.cassandra.exceptions.QueryCancelledException;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.disk.SSTableRowIdKeyRangeIterator;
+import org.apache.cassandra.index.sai.disk.v1.postings.MergePostingList;
+import org.apache.cassandra.index.sai.disk.v1.vector.CheckpointingIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator;
+import org.apache.cassandra.index.sai.metrics.TableQueryMetrics;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.plan.StorageAttachedIndexQueryController;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.postings.PostingListIntersection;
+import org.apache.cassandra.index.sai.postings.PostingListUnion;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.utils.Throwables;
+
+public class VectorQueryController extends StorageAttachedIndexQueryController
+{
+ private static final Logger logger = LoggerFactory.getLogger(VectorQueryController.class);
+
+ public VectorQueryController(ColumnFamilyStore cfs,
+ ReadCommand command,
+ RowFilter filterOperation,
+ QueryContext queryContext,
+ TableQueryMetrics tableQueryMetrics)
+ {
+ super(cfs, command, filterOperation, queryContext, tableQueryMetrics);
+ }
+
+ /**
+ * Build a {@link KeyRangeIterator.Builder} from the given list of expressions by applying given operation (OR/AND).
+ * Building of such builder involves index search, results of which are persisted in the internal resources list
+ *
+ * @param expressions The expressions to build range iterator from (expressions with not results are ignored).
+ *
+ * @return range iterator builder based on given expressions and operation type.
+ */
+ public KeyRangeIterator.Builder getIndexQueryResults(Collection<Expression> expressions)
+ {
+ // FIXME having this at the expression level means that it only gets applied to Nodes at that level;
+ // moving it to ORDER BY will allow us to apply it correctly for other Node sub-trees
+ var annExpressionInHybridSearch = getAnnExpressionInHybridSearch(expressions);
+ boolean isAnnHybridSearch = annExpressionInHybridSearch != null;
+ if (isAnnHybridSearch)
+ expressions = expressions.stream().filter(e -> e != annExpressionInHybridSearch).collect(Collectors.toList());
+
+ // search memtable before referencing sstable indexes; otherwise we may miss newly flushed memtable index
+ var iteratorsByMemtable = expressions.stream()
+ .filter(expr -> !expr.context.isNotIndexed())
+ .flatMap(expr -> expr.context.getMemtableIndexManager().iteratorsForSearch(queryContext, expr, mergeRange).stream())
+ .collect(Collectors.groupingBy(pair -> pair.left, Collectors.mapping(pair -> pair.right, Collectors.toList())));
+
+ VectorQueryViewBuilder.QueryView queryView = new VectorQueryViewBuilder(expressions, mergeRange).build();
+ // in case of ANN query in hybrid search, we have to reference ANN sstable indexes separately because queryView doesn't include ANN sstable indexes
+ VectorQueryViewBuilder.QueryView annQueryViewInHybridSearch = isAnnHybridSearch ? new VectorQueryViewBuilder(Collections.singleton(annExpressionInHybridSearch), mergeRange).build() : null;
+
+ try
+ {
+ List<KeyRangeIterator> sstableIntersections = new ArrayList<>(queryView.view.size());
+ for (Map.Entry<SSTableReader, List<VectorQueryViewBuilder.IndexExpression>> viewEntry : queryView.view.entrySet())
+ {
+ PostingList postings = createRowIdIterator(viewEntry.getValue(), isAnnHybridSearch);
+ if (isAnnHybridSearch)
+ sstableIntersections.add(reorderAndLimitBySSTableRowIds(postings, viewEntry.getKey(), annQueryViewInHybridSearch));
+ else
+ {
+ PrimaryKeyMap.Factory pkFactory = viewEntry.getValue().iterator().next().index.getSSTableContext().primaryKeyMapFactory;
+ sstableIntersections.add(convertToPrimaryKeyIterator(pkFactory, postings));
+ }
+ }
+
+ List<KeyRangeIterator> memtableIntersections = iteratorsByMemtable.entrySet()
+ .stream()
+ .map(e -> {
+ // we need to do all the intersections at the index level, or ordering won't work
+ KeyRangeIterator it = buildIterator(e.getValue(), isAnnHybridSearch);
+ if (isAnnHybridSearch)
+ it = reorderAndLimitBy(it, e.getKey(), annExpressionInHybridSearch);
+ return it;
+ })
+ .collect(Collectors.toList());
+
+ Iterable<KeyRangeIterator> allIntersections = Iterables.concat(sstableIntersections, memtableIntersections);
Review Comment:
This is the commit I'm intending bringing in. I believe this will address a number of issues.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358470397
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -281,14 +288,29 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
Selectors selectors = selection.newSelectors(options);
AggregationSpecification aggregationSpec = getAggregationSpec(options);
- ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
- nowInSec, userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+ DataLimits limit = getDataLimits(userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+
+ // Handle additional validation for topK queries
+ if (restrictions.isTopK())
+ {
+ checkFalse(aggregationSpec != null, TOPK_AGGREGATION_ERROR);
+
+ checkFalse(options.getConsistency() != ConsistencyLevel.ONE &&
Review Comment:
I have pushed a commit that, hopefully, meets both of these requirements.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1337139772
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -483,7 +486,7 @@ else if (restrictions.keyIsInRelation())
// Please note that the isExhausted state of the pager only gets updated when we've closed the page, so this
// shouldn't be moved inside the 'try' above.
- if (!pager.isExhausted())
+ if (!pager.isExhausted() && !pager.pager.isTopK())
Review Comment:
The query issues were a 2 part problem in the `VectorQueryController`.
1. The code to establish whether a query was a hybrid query or not did not consider whether the other expressions in the query were indexed or not.
2. The code to build the iterators was not checking for indexed expressions correctly.
The `UOE` on flush is not reproducing locally, and the test with flush is running correctly. I can only assume that this has been fixed elsewhere.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1333310626
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -535,7 +538,7 @@ public ResultMessage.Rows executeInternal(QueryState state,
try (ReadExecutionController executionController = query.executionController())
{
- if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize)))
+ if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize) || query.isTopK()))
Review Comment:
It seems that using ANN on an aggregation query will start paging (that we don't support). The query will be rejected by SAI validation as requiring a page limit, even if it actually has one:
```java
@Test
public void testAggregation()
{
createTable("CREATE TABLE %s (k int PRIMARY KEY, v vector<float, 1>, c int)");
createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
execute("INSERT INTO %s (k, v, c) VALUES (1, [4], 1)");
execute("INSERT INTO %s (k, v, c) VALUES (2, [3], 10)");
execute("INSERT INTO %s (k, v, c) VALUES (3, [2], 100)");
execute("INSERT INTO %s (k, v, c) VALUES (4, [1], 1000)");
assertRows(execute("SELECT sum(c) FROM %s"), row(1111));
assertRows(execute("SELECT sum(c), v FROM %s"), row(1111, vector(4f)));
assertRows(execute("SELECT sum(c) FROM %s WHERE k = 1 ORDER BY v ANN OF [0] LIMIT 4"), row(1111, vector(1f))); // rejected!
}
```
Should we forbid aggregation queries with ANN?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1340216519
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1133,13 +1149,13 @@ public SelectStatement prepare(ClientState state, boolean forView) throws Invali
List<Selectable> selectables = RawSelector.toSelectables(selectClause, table);
boolean containsOnlyStaticColumns = selectOnlyStaticColumns(table, selectables);
- StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, containsOnlyStaticColumns, forView);
+ List<Ordering> orderings = getOrderings(table);
+ StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, orderings, containsOnlyStaticColumns, forView);
// If we order post-query, the sorted column needs to be in the ResultSet for sorting,
// even if we don't ultimately ship them to the client (CASSANDRA-4911).
- Map<ColumnMetadata, Boolean> orderingColumns = getOrderingColumns(table);
- Set<ColumnMetadata> resultSetOrderingColumns = restrictions.keyIsInRelation() ? orderingColumns.keySet()
- : Collections.emptySet();
+ Map<ColumnMetadata, Ordering> orderingColumns = getOrderingColumns(orderings);
+ Set<ColumnMetadata> resultSetOrderingColumns = getResultSetOrdering(restrictions, orderingColumns);
Review Comment:
This failure is reproducible without SAI involved at all. The following test also fails on the `cassandra-5.0` branch:
```
@Test
public void test2() throws IOException
{
try (Cluster cluster = init(Cluster.build(1).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int, c int, v int, primary key(k, c))"));
cluster.coordinator(1).execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (0, 1, 2)"), ConsistencyLevel.QUORUM);
String query = withKeyspace("SELECT v FROM %s.t WHERE k IN (0, 1) ORDER BY c LIMIT 10");
assertRows(cluster.coordinator(1).execute(query, ConsistencyLevel.ONE), row(2));
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1340232886
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1133,13 +1149,13 @@ public SelectStatement prepare(ClientState state, boolean forView) throws Invali
List<Selectable> selectables = RawSelector.toSelectables(selectClause, table);
boolean containsOnlyStaticColumns = selectOnlyStaticColumns(table, selectables);
- StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, containsOnlyStaticColumns, forView);
+ List<Ordering> orderings = getOrderings(table);
+ StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, orderings, containsOnlyStaticColumns, forView);
// If we order post-query, the sorted column needs to be in the ResultSet for sorting,
// even if we don't ultimately ship them to the client (CASSANDRA-4911).
- Map<ColumnMetadata, Boolean> orderingColumns = getOrderingColumns(table);
- Set<ColumnMetadata> resultSetOrderingColumns = restrictions.keyIsInRelation() ? orderingColumns.keySet()
- : Collections.emptySet();
+ Map<ColumnMetadata, Ordering> orderingColumns = getOrderingColumns(orderings);
+ Set<ColumnMetadata> resultSetOrderingColumns = getResultSetOrdering(restrictions, orderingColumns);
Review Comment:
I think both of these make sense because without the ordering column being returned it would be impossible to determine the ordering logic. If the above query is changed to `SELECT v FROM %s.t WHERE k = 0 ORDER BY c LIMIT 10`, then the clustering key is not returned.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342846250
##########
src/java/org/apache/cassandra/index/sai/IndexContext.java:
##########
@@ -354,6 +371,11 @@ public boolean isLiteral()
return TypeUtil.isLiteral(getValidator());
}
+ public boolean isVector()
+ {
+ return getValidator().isVector() && ((VectorType<?>)getValidator()).elementType instanceof FloatType;
Review Comment:
At the moment if an index is created on a vector column the type of the created index will depend on the type of the vector elements. If the vector is made of floats we will create and ANN index that doesn't support equality restrictions. If the vector is made of any other thing we wil create a literal index that does support equalities.
I think this might be a bit confussing for users. They will be able to build indexes for equalities on any kind of vectors but float vectors. I guess we could:
1. Forbid SAI indexes in non-float vectors, so vectors are always indexed by ANN if ANN supports their element type (float)
2. Vectors are indexed as literals by default, unless some syntax or index property indicates that they should be indexed by ANN.
wdyt?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341284886
##########
src/java/org/apache/cassandra/index/Index.java:
##########
@@ -440,6 +445,19 @@ default boolean filtersMultipleContains()
*/
public RowFilter getPostIndexQueryFilter(RowFilter filter);
+ /**
+ * Return a comparator that reorders query result before sending to client
+ *
+ * @param restriction restriction that requires current index
+ * @param columnIndex idx of the indexed column in returned row
+ * @param options query options
+ * @return comparator that for post-query ordering; or null if not supported
+ */
+ default Comparator<List<ByteBuffer>> getPostQueryOrdering(Restriction restriction, int columnIndex, QueryOptions options)
Review Comment:
Ok, I took a bit of time to understand what you were asking but this does make more sense because it disconnects the comparator from the row. I have made these changes
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1337053163
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1573,6 +1656,39 @@ public int compare(List<ByteBuffer> a, List<ByteBuffer> b)
}
}
+ private static class IndexColumnComparator<T> extends ColumnComparator<List<ByteBuffer>>
+ {
+ private final SingleRestriction restriction;
+ private final int columnIndex;
+
+ // VSTODO maybe cache in prepared statement
+ public IndexColumnComparator(SingleRestriction restriction, int columnIndex)
+ {
+ this.restriction = restriction;
+ this.columnIndex = columnIndex;
+ }
+
+ @Override
+ public boolean indexOrdering()
+ {
+ return true;
+ }
+
+ @Override
+ public Comparator<List<ByteBuffer>> prepareFor(TableMetadata table, QueryOptions options)
+ {
+ Index index = restriction.findSupportingIndex(IndexRegistry.obtain(table));
Review Comment:
I agree. I have added a `findSupportingIndexFromQueryPlan` to `Restriction` and am using that here instead. While it is unlikely that this is a problem, it is a good idea to have the index lookups consistent.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357044900
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -476,6 +507,38 @@ public boolean usesSecondaryIndexing()
return this.usesSecondaryIndexing;
}
+ /**
+ * This is a hack to push ordering down to indexes.
+ * Indexes are selected based on RowFilter only, so we need to turn orderings into restrictions
+ * so they end up in the row filter.
+ *
+ * @param orderings orderings from the select statement
+ * @return the {@link RestrictionSet} with the added orderings
+ */
+ private RestrictionSet addOrderingRestrictions(List<Ordering> orderings, RestrictionSet restrictionSet)
+ {
+ List<Ordering> annOrderings = orderings.stream().filter(o -> o.expression.hasNonClusteredOrdering()).collect(Collectors.toList());
+
+ if (annOrderings.size() > 1)
+ throw new InvalidRequestException("Cannot specify more than one ANN ordering");
+ else if (annOrderings.size() == 1)
+ {
+ if (orderings.size() > 1)
+ throw new InvalidRequestException("ANN ordering does not support secondary ordering");
Review Comment:
+1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357158520
##########
src/java/org/apache/cassandra/index/sai/plan/vector/VectorQueryController.java:
##########
@@ -0,0 +1,247 @@
+/*
+ * 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.index.sai.plan.vector;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.filter.RowFilter;
+import org.apache.cassandra.db.memtable.Memtable;
+import org.apache.cassandra.exceptions.QueryCancelledException;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.disk.SSTableRowIdKeyRangeIterator;
+import org.apache.cassandra.index.sai.disk.v1.postings.MergePostingList;
+import org.apache.cassandra.index.sai.disk.v1.vector.CheckpointingIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator;
+import org.apache.cassandra.index.sai.metrics.TableQueryMetrics;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.plan.StorageAttachedIndexQueryController;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+import org.apache.cassandra.index.sai.postings.PostingList;
+import org.apache.cassandra.index.sai.postings.PostingListIntersection;
+import org.apache.cassandra.index.sai.postings.PostingListUnion;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.utils.Throwables;
+
+public class VectorQueryController extends StorageAttachedIndexQueryController
Review Comment:
i'm worried that having a bunch of subclasses for these is going to make maintenance more difficult
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357148219
##########
src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java:
##########
@@ -184,4 +179,43 @@ private long flush(PrimaryKey minKey,
return numRows;
}
+
+ private void flushVectorIndex(long startTime, Stopwatch stopwatch) throws IOException
+ {
+ SegmentMetadata.ComponentMetadataMap metadataMap = memtable.writeDirect(indexDescriptor, indexContext, rowMapping::get);
+
+ completeIndexFlush(rowMapping.size(), startTime, stopwatch);
+
+ SegmentMetadata metadata = new SegmentMetadata(0,
+ rowMapping.size(),
+ 0,
+ rowMapping.maxSSTableRowId,
+ rowMapping.minKey,
+ rowMapping.maxKey,
+ ByteBufferUtil.bytes(0), // VSTODO by pass min max terms for vectors
Review Comment:
are these TODOs still relevant?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357060175
##########
src/java/org/apache/cassandra/db/marshal/ByteArrayAccessor.java:
##########
@@ -236,6 +236,19 @@ public float toFloat(byte[] value)
return ByteArrayUtil.getFloat(value, 0);
}
+ @Override
+ public float[] toFloatArray(byte[] value, int dimension)
Review Comment:
for large vectors it is probably going to be faster to leverage FloatBuffer.get, but that can be a future optimization
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358486563
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java:
##########
@@ -0,0 +1,56 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+
+/**
+ * There are two steps in ordering:
Review Comment:
Which bit is outdated? The first block is the same as on the vsearch branch. The second block we added to try and document the second step.
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java:
##########
@@ -0,0 +1,56 @@
+/*
+ * 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.index.sai.disk.v1.segment;
+
+import java.io.IOException;
+
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.db.partitions.PartitionIterator;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.plan.Expression;
+import org.apache.cassandra.index.sai.postings.PeekablePostingList;
+
+/**
+ * There are two steps in ordering:
+ * <p>
+ * 1. Limit a single sstable's results to the correct keys. At this stage
+ * we put them back in primary key order to play nice with the rest of
+ * the query pipeline.
+ * 2. Merge the results from multiple sstables. Now we leave them in the
+ * final, correct order.
+ * <p>
+ * SegmentOrdering handles the first step.
+ *
+ * The second step is handled during post-reconciliation processing when the
+ * {@link org.apache.cassandra.index.Index.QueryPlan#postProcessor(ReadCommand)} is called by
+ * {@link org.apache.cassandra.db.PartitionRangeReadCommand#postReconciliationProcessing(PartitionIterator)}. If the
+ * {@link org.apache.cassandra.index.Index.QueryPlan} is for a topK query then a reordering filter is returned to return
+ * the partitions in the correct order.
+ */
+public interface SegmentOrdering
+{
+ /**
+ * Reorder, limit, and put back into original order the results from a single sstable
Review Comment:
As above.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1360586490
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -109,6 +112,15 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(SelectStatement.logger, 1, TimeUnit.MINUTES);
public static final int DEFAULT_PAGE_SIZE = 10000;
+ public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top-K queries can only be run with consistency level ONE/LOCAL_ONE. Consistency level %s was used.";
+ public static final String TOPK_LIMIT_ERROR = "Top-K queries must have a limit specified and the limit must be less than the query page size";
+ public static final String TOPK_PARTITION_LIMIT_ERROR = "Top-K queries do not support per-partition limits";
+ public static final String TOPK_AGGREGATION_ERROR = "Top-K queries can not be run with aggregation";
+ public static final String TOPK_CONSISTENCY_LEVEL_WARNING = "Top-K queries can only be run with consistency level ONE " +
+ "/ LOCAL_ONE / NODE_LOCAL. Consistency level %s was used. " +
+ "Downgrading to consistency level %s.";
Review Comment:
The message says that top-k queries can only be run with ONE/LOCAL_ONE/NODE_LOCAL. However, the code at `execute` will rewrite any code distinct to ONE/LOCAL_ONE to be transformed into either ONE or LOCAL_ONE. So `NODE_LOCAL` will always be rewritten to `LOCAL_ONE`, producing the confusing message `Top-K queries can only be run with consistency level ONE / LOCAL_ONE / NODE_LOCAL. Consistency level NODE_LOCAL was used. Downgrading to consistency level LOCAL_ONE.`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1333005479
##########
src/java/org/apache/cassandra/index/Index.java:
##########
@@ -440,6 +445,19 @@ default boolean filtersMultipleContains()
*/
public RowFilter getPostIndexQueryFilter(RowFilter filter);
+ /**
+ * Return a comparator that reorders query result before sending to client
+ *
+ * @param restriction restriction that requires current index
+ * @param columnIndex idx of the indexed column in returned row
+ * @param options query options
+ * @return comparator that for post-query ordering; or null if not supported
+ */
+ default Comparator<List<ByteBuffer>> getPostQueryOrdering(Restriction restriction, int columnIndex, QueryOptions options)
Review Comment:
Ah, it only takes a restriction, so it's fixed to a single-column ordering. In that case, I think we could omit the column index arg and make it return a column value comparator, instead of a row comparator:
```suggestion
default Comparator<ByteBuffer> getPostQueryOrdering(Restriction restriction, QueryOptions options)
```
So the CQL layer could take care of dealing with the column index in the result set. For example, in `SelectStatement.IndexColumnComparator#prepareFor`:
```java
Comparator<ByteBuffer> comparator = index.getPostQueryOrdering(restriction, options);
return (a, b) -> compare(comparator, a.get(columnIndex), b.get(columnIndex));
```
And the SAI implementation would be:
```java
@Override
public Comparator<ByteBuffer> getPostQueryOrdering(Restriction restriction, QueryOptions options)
{
// For now, only support ANN
assert restriction instanceof SingleColumnRestriction.AnnRestriction;
Preconditions.checkState(indexContext.isVector());
SingleColumnRestriction.AnnRestriction annRestriction = (SingleColumnRestriction.AnnRestriction) restriction;
VectorSimilarityFunction function = indexContext.getIndexWriterConfig().getSimilarityFunction();
float[] target = TypeUtil.decomposeVector(indexContext, annRestriction.value(options).duplicate());
return (leftBuf, rightBuf) -> {
float[] left = TypeUtil.decomposeVector(indexContext, leftBuf.duplicate());
double scoreLeft = function.compare(left, target);
float[] right = TypeUtil.decomposeVector(indexContext, rightBuf.duplicate());
double scoreRight = function.compare(right, target);
return Double.compare(scoreRight, scoreLeft); // descending order
};
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1363984958
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
Review Comment:
Shouldn't this be mentioned in the list of limitations on the usage warning?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327696483
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -135,6 +135,9 @@ private boolean needsReplicaFilteringProtection()
if (command.rowFilter().isEmpty())
return false;
+ if (command.isTopK())
+ return false;
Review Comment:
Here is a dtest showing how a vector search with CL>ONE can wrongly return a row that has been deleted:
```java
try (Cluster cluster = init(Cluster.build(2).withConfig(config -> config.with(GOSSIP, NETWORK)).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int PRIMARY KEY, v vector<float, 1>)"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex'"));
SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
// insert a row in both nodes, and then remove in only one of the nodes
cluster.coordinator(1).execute(withKeyspace("INSERT INTO %s.t (k, v) VALUES (1, [1])"), ConsistencyLevel.QUORUM);
cluster.get(1).executeInternal(withKeyspace("DELETE FROM %s.t WHERE k=1"));
String query = withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [1] LIMIT 10");
Object[][] rows = cluster.coordinator(1).execute(query, ConsistencyLevel.QUORUM);
assertThat(rows).isEmpty(); // wrongly returns the deleted row
}
```
This other
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327723617
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -135,6 +135,9 @@ private boolean needsReplicaFilteringProtection()
if (command.rowFilter().isEmpty())
return false;
+ if (command.isTopK())
+ return false;
Review Comment:
And, since RFP is disabled for all the query, if we restrict any other column we can also return stale rows that don't satisfy that filter:
```java
try (Cluster cluster = init(Cluster.build(2).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int PRIMARY KEY, v vector<float, 1>, c text)"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex'"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(c) USING 'StorageAttachedIndex'"));
// insert a row in both nodes, and then update it in only one of the nodes
cluster.coordinator(1).execute(withKeyspace("INSERT INTO %s.t (k, v, c) VALUES (1, [1], 'old')"), ConsistencyLevel.QUORUM);
cluster.get(1).executeInternal(withKeyspace("UPDATE %s.t SET c='new' WHERE k=1"));
String query = withKeyspace("SELECT * FROM %s.t WHERE c='old' ORDER BY v ANN OF [1] LIMIT 1");
Object[][] rows = cluster.coordinator(1).execute(query, ConsistencyLevel.QUORUM);
assertThat(rows).isEmpty(); // wrongly returns the stale row, (1, old, [1.0]), which doesn't satisfy the query
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327951896
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1133,13 +1149,13 @@ public SelectStatement prepare(ClientState state, boolean forView) throws Invali
List<Selectable> selectables = RawSelector.toSelectables(selectClause, table);
boolean containsOnlyStaticColumns = selectOnlyStaticColumns(table, selectables);
- StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, containsOnlyStaticColumns, forView);
+ List<Ordering> orderings = getOrderings(table);
+ StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, orderings, containsOnlyStaticColumns, forView);
// If we order post-query, the sorted column needs to be in the ResultSet for sorting,
// even if we don't ultimately ship them to the client (CASSANDRA-4911).
- Map<ColumnMetadata, Boolean> orderingColumns = getOrderingColumns(table);
- Set<ColumnMetadata> resultSetOrderingColumns = restrictions.keyIsInRelation() ? orderingColumns.keySet()
- : Collections.emptySet();
+ Map<ColumnMetadata, Ordering> orderingColumns = getOrderingColumns(orderings);
+ Set<ColumnMetadata> resultSetOrderingColumns = getResultSetOrdering(restrictions, orderingColumns);
Review Comment:
It seems that somehow the vector column of an ANN query is always included in the results set, even if it's not included in the selection. This single-node dtest shows the issue:
```java
try (Cluster cluster = init(Cluster.build(1).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int PRIMARY KEY, v vector<float, 1>, c int)"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex'"));
cluster.coordinator(1).execute(withKeyspace("INSERT INTO %s.t (k, v, c) VALUES (0, [1], 2)"), ConsistencyLevel.QUORUM);
String query = withKeyspace("SELECT c FROM %s.t ORDER BY v ANN OF [1] LIMIT 10");
assertRows(cluster.coordinator(1).execute(query, ConsistencyLevel.QUORUM), row(2)); // includes the vector column!!
}
```
We can also reproduce it with a unit test, but it requires to define a page size, as I think drivers usually do:
```java
@Test
public void testSelection() throws Throwable
{
createTable("CREATE TABLE %s (k int PRIMARY KEY, v vector<float, 1>, c int)");
createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
execute("INSERT INTO %s (k, v, c) VALUES (0, [1], 2)");
String query = "SELECT c FROM %s ORDER BY v ANN OF [1] LIMIT 10";
assertRows(execute(query), row(2));
requireNetwork();
assertRowsNet(executeNet(query), row(2));
assertRowsNet(executeNetWithPaging(query, 1000), row(2)); // wrongly includes the vector column
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1368551933
##########
src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java:
##########
@@ -184,4 +179,43 @@ private long flush(PrimaryKey minKey,
return numRows;
}
+
+ private void flushVectorIndex(long startTime, Stopwatch stopwatch) throws IOException
+ {
+ SegmentMetadata.ComponentMetadataMap metadataMap = memtable.writeDirect(indexDescriptor, indexContext, rowMapping::get);
+
+ completeIndexFlush(rowMapping.size(), startTime, stopwatch);
+
+ SegmentMetadata metadata = new SegmentMetadata(0,
+ rowMapping.size(),
+ 0,
+ rowMapping.maxSSTableRowId,
+ rowMapping.minKey,
+ rowMapping.maxKey,
+ ByteBufferUtil.bytes(0), // VSTODO by pass min max terms for vectors
Review Comment:
It isn't really a TODO any more but we could eventually tidy this so we could just not pass anything.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357067949
##########
src/java/org/apache/cassandra/index/sai/disk/format/IndexComponent.java:
##########
@@ -50,6 +50,11 @@ public enum IndexComponent
*/
TERMS_DATA("TermsData"),
+ /**
+ * Product Quantization store used to store compressed vectors for the vector index
+ */
+ VECTORS("Vectors"),
Review Comment:
The javadoc is correct but IMO the name is misleading since the actual vectors are stored inline with the graph.
COMPRESSED_VECTORS is a bit of a mouthful, PQ_VECTORS is probably more reasonable
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357040047
##########
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java:
##########
@@ -424,8 +424,9 @@ public enum CassandraRelevantProperties
SAI_INTERSECTION_CLAUSE_LIMIT("cassandra.sai.intersection_clause_limit", "2"),
/** Latest version to be used for SAI index writing */
SAI_LATEST_VERSION("cassandra.sai.latest_version", "aa"),
- SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size_kb", "5"),
- SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size_kb", "1"),
+ SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size", "5KiB"),
+ SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size", "1KiB"),
+ SAI_MAX_VECTOR_TERM_SIZE("cassandra.sai.max_vector_term_size", "16KiB"),
Review Comment:
Besides rising the default, perhaps we might want to express the term size as the number of dimensions, given that we only support vectors of floats? Like in `cassandra.sai.max_vector_dimensions`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356778159
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java:
##########
@@ -0,0 +1,673 @@
+/*
+ * 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.index.sai.cql;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.junit.Test;
+
+import org.apache.cassandra.config.CassandraRelevantProperties;
+import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.service.ClientWarn;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class VectorTypeTest extends VectorTester
+{
+ private static final IPartitioner partitioner = Murmur3Partitioner.instance;
+
+ @Test
+ public void endToEndTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', [5.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (5, 'F', [6.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (6, 'G', [7.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (7, 'H', [8.0, 5.0, 6.0])");
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5");
+ assertThat(result).hasSize(5);
+
+ // some data that only lives in memtable
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (8, 'I', [9.0, 5.0, 6.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (9, 'J', [10.0, 6.0, 7.0])");
+ result = execute("SELECT * FROM %s ORDER BY val ann of [9.5, 5.5, 6.5] LIMIT 5");
+ assertContainsInt(result, "pk", 8);
+ assertContainsInt(result, "pk", 9);
+
+ // data from sstables
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertContainsInt(result, "pk", 1);
+ assertContainsInt(result, "pk", 2);
+ }
+
+ @Test
+ public void warningIsIssuedOnIndexCreation()
+ {
+ ClientWarn.instance.captureWarnings();
+ createTable("CREATE TABLE %s (pk int, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+ List<String> warnings = ClientWarn.instance.getWarnings();
+
+ assertTrue(warnings.size() > 0);
+ assertEquals(StorageAttachedIndex.VECTOR_USAGE_WARNING, warnings.get(0));
+ }
+
+ @Test
+ public void createIndexAfterInsertTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ flush();
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+ }
+
+ public static void assertContainsInt(UntypedResultSet result, String columnName, int columnValue)
+ {
+ for (UntypedResultSet.Row row : result)
+ {
+ if (row.has(columnName))
+ {
+ int value = row.getInt(columnName);
+ if (value == columnValue)
+ {
+ return;
+ }
+ }
+ }
+ throw new AssertionError("Result set does not contain a row with " + columnName + " = " + columnValue);
+ }
+
+ @Test
+ public void testTwoPredicates()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, b, v) VALUES (0, true, [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (1, true, [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (2, false, [3.0, 4.0, 5.0])");
+
+ // the vector given is closest to row 2, but we exclude that row because b=false
+ var result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ // VSTODO assert specific row keys
+ assertThat(result).hasSize(2);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testTwoPredicatesManyRows()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ for (int i = 0; i < 100; i++)
+ execute("INSERT INTO %s (pk, b, v) VALUES (?, true, ?)",
+ i, vector((float) i, (float) (i + 1), (float) (i + 2)));
+
+ var result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testThreePredicates()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, str text, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(str) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (0, true, [1.0, 2.0, 3.0], 'A')");
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (1, true, [2.0, 3.0, 4.0], 'B')");
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (2, false, [3.0, 4.0, 5.0], 'C')");
+
+ // the vector given is closest to row 2, but we exclude that row because b=false and str!='B'
+ var result = execute("SELECT * FROM %s WHERE b=true AND str='B' ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ // VSTODO assert specific row keys
+ assertThat(result).hasSize(1);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true AND str='B' ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void testSameVectorMultipleRows()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'A', [1.0, 2.0, 3.0])");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+ }
+
+ @Test
+ public void testQueryEmptyTable()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertThat(result).hasSize(0);
+ }
+
+ @Test
+ public void testQueryTableWithNulls()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', null)");
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertThat(result).hasSize(0);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [4.0, 5.0, 6.0])");
+ result = execute("SELECT pk FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertRows(result, row(1));
+ }
+
+ @Test
+ public void testLimitLessThanInsertedRowCount()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ // Insert more rows than the query limit
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [4.0, 5.0, 6.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [7.0, 8.0, 9.0])");
+
+ // Query with limit less than inserted row count
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testQueryMoreRowsThanInserted()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void changingOptionsTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ if (CassandraRelevantProperties.SAI_VECTOR_ALLOW_CUSTOM_PARAMETERS.getBoolean())
+ {
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = " +
+ "{'maximum_node_connections' : 10, 'construction_beam_width' : 200, 'similarity_function' : 'euclidean' }");
+ }
+ else
+ {
+ assertThatThrownBy(() -> createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = " +
+ "{'maximum_node_connections' : 10, 'construction_beam_width' : 200, 'similarity_function' : 'euclidean' }"))
+ .isInstanceOf(InvalidRequestException.class);
+ return;
+ }
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', [5.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (5, 'F', [6.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (6, 'G', [7.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (7, 'H', [8.0, 5.0, 6.0])");
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5");
+ assertThat(result).hasSize(5);
+ }
+
+ @Test
+ public void bindVariablesTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', ?)", vector(2.0f ,3.0f, 4.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', ?)", vector(4.0f, 5.0f, 6.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of ? LIMIT 3", vector(2.5f, 3.5f, 4.5f));
+ assertThat(result).hasSize(3);
+ }
+
+ @Test
+ public void intersectedSearcherTest()
+ {
+ // check that we correctly get back the two rows with str_val=B even when those are not
+ // the closest rows to the query vector
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', ?)", vector(2.0f ,3.0f, 4.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'B', ?)", vector(4.0f, 5.0f, 6.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', ?)", vector(5.0f, 6.0f, 7.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+
+ flush();
+ result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void nullVectorTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val) VALUES (1, 'B')"); // no vector
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val) VALUES (3, 'D')"); // no vector
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', ?)", vector(5.0f, 6.0f, 7.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(0);
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'A' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+
+ flush();
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(0);
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'A' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void lwtTest()
+ {
+ createTable("CREATE TABLE %s (p int, c int, v text, vec vector<float, 2>, PRIMARY KEY(p, c))");
+ createIndex("CREATE CUSTOM INDEX ON %s(vec) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (p, c, v) VALUES (?, ?, ?)", 0, 0, "test");
+ execute("INSERT INTO %s (p, c, v) VALUES (?, ?, ?)", 0, 1, "00112233445566");
+
+ execute("UPDATE %s SET v='00112233', vec=[0.9, 0.7] WHERE p = 0 AND c = 0 IF v = 'test'");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY vec ANN OF [0.1, 0.9] LIMIT 100");
+
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void twoVectorFieldsTest()
+ {
+ createTable("CREATE TABLE %s (pk int, v2 vector<float, 2>, v3 vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(v2) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v3) USING 'StorageAttachedIndex'");
+ }
+
+ @Test
+ public void primaryKeySearchTest()
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, 3>, i int, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ var N = 5;
+ for (int i = 0; i < N; i++)
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", i, vector(1.0f + i, 2.0f + i, 3.0f + i));
+
+ for (int i = 0; i < N; i++)
+ {
+ UntypedResultSet result = execute("SELECT pk FROM %s WHERE pk = ? ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2", i);
+ assertThat(result).hasSize(1);
+ assertRows(result, row(i));
+ }
+
+ flush();
+ for (int i = 0; i < N; i++)
+ {
+ UntypedResultSet result = execute("SELECT pk FROM %s WHERE pk = ? ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2", i);
+ assertThat(result).hasSize(1);
+ assertRows(result, row(i));
+ }
+ }
+
+ @Test
+ public void partitionKeySearchTest()
+ {
+ createTable("CREATE TABLE %s (partition int, row int, val vector<float, 2>, PRIMARY KEY(partition, row))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
+
+ var nPartitions = 5;
+ var rowsPerPartition = 10;
+ Map<Integer, List<float[]>> vectorsByPartition = new HashMap<>();
+
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ for (int j = 1; j <= rowsPerPartition; j++)
+ {
+ logger.debug("Inserting partition {} row {}: [{}, {}]", i, j, i, j);
+ execute("INSERT INTO %s (partition, row, val) VALUES (?, ?, ?)", i, j, vector((float) i, (float) j));
+ float[] vector = {(float) i, (float) j};
+ vectorsByPartition.computeIfAbsent(i, k -> new ArrayList<>()).add(vector);
+ }
+ }
+
+ var queryVector = vector(new float[] { 1.5f, 1.5f });
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition, row FROM %s WHERE partition = ? ORDER BY val ann of ? LIMIT 2", i, queryVector);
+ assertThat(result).hasSize(2);
+ assertRowsIgnoringOrder(result,
+ row(i, 1),
+ row(i, 2));
+ }
+
+ flush();
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition, row FROM %s WHERE partition = ? ORDER BY val ann of ? LIMIT 2", i, queryVector);
+ assertThat(result).hasSize(2);
+ assertRowsIgnoringOrder(result,
+ row(i, 1),
+ row(i, 2));
+ }
+ }
+
+ @Test
+ public void clusteringKeyIndexTest()
+ {
+ createTable("CREATE TABLE %s (pk int, ck vector<float, 2>, PRIMARY KEY(pk, ck))");
+ createIndex("CREATE CUSTOM INDEX ON %s(ck) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, ck) VALUES (1, [1.0, 2.0])");
+
+ assertRows(execute("SELECT * FROM %s ORDER BY ck ANN OF [1.0, 2.0] LIMIT 1"), row(1, vector(1.0F, 2.0F)));
+ }
+
+ @Test
+ public void rangeSearchTest() throws Throwable
+ {
+ createTable("CREATE TABLE %s (partition int, val vector<float, 2>, PRIMARY KEY(partition))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
+
+ var nPartitions = 100;
+ Map<Integer, float[]> vectorsByKey = new HashMap<>();
+
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ float[] vector = {(float) i, (float) i};
+ execute("INSERT INTO %s (partition, val) VALUES (?, ?)", i, vector(vector));
+ vectorsByKey.put(i, vector);
+ }
+
+ var queryVector = vector(new float[] { 1.5f, 1.5f });
+ CheckedFunction tester = () -> {
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithLowerBound(vectorsByKey.keySet(), i, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithLowerBound(vectorsByKey.keySet(), i, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithUpperBound(vectorsByKey.keySet(), i, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithUpperBound(vectorsByKey.keySet(), i, true));
+
+ for (int j = 1; j <= nPartitions; j++)
+ {
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) AND token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, true, j, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) AND token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, false, j, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) AND token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, true, j, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) AND token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, false, j, false));
+ }
+ }
+ };
+
+ tester.apply();
+
+ flush();
+
+ tester.apply();
+ }
+
+ private Collection<Integer> keys(UntypedResultSet result)
+ {
+ List<Integer> keys = new ArrayList<>(result.size());
+ for (UntypedResultSet.Row row : result)
+ keys.add(row.getInt("partition"));
+ return keys;
+ }
+
+ private Collection<Integer> keysWithLowerBound(Collection<Integer> keys, int leftKey, boolean leftInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getToken(Int32Type.instance.decompose(leftKey)), leftInclusive,
+ partitioner.getMaximumToken().getToken(), true);
+ }
+
+ private Collection<Integer> keysWithUpperBound(Collection<Integer> keys, int rightKey, boolean rightInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getMinimumToken().getToken(), true,
+ partitioner.getToken(Int32Type.instance.decompose(rightKey)), rightInclusive);
+ }
+
+ private Collection<Integer> keysInBounds(Collection<Integer> keys, int leftKey, boolean leftInclusive, int rightKey, boolean rightInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getToken(Int32Type.instance.decompose(leftKey)), leftInclusive,
+ partitioner.getToken(Int32Type.instance.decompose(rightKey)), rightInclusive);
+ }
+
+ private Collection<Integer> keysInTokenRange(Collection<Integer> keys, Token leftToken, boolean leftInclusive, Token rightToken, boolean rightInclusive)
+ {
+ long left = leftToken.getLongValue();
+ long right = rightToken.getLongValue();
+ return keys.stream()
+ .filter(k -> {
+ long t = partitioner.getToken(Int32Type.instance.decompose(k)).getLongValue();
+ return (left < t || left == t && leftInclusive) && (t < right || t == right && rightInclusive);
+ }).collect(Collectors.toSet());
+ }
+
+ @Test
+ public void selectFloatVectorFunctions()
+ {
+ createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, value vector<float, 2>)");
+
+ // basic functionality
+ Vector<Float> q = vector(1f, 2f);
+ execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(1f, 2f));
+ execute("SELECT similarity_cosine(value, value) FROM %s WHERE pk=0");
+
+ // type inference checks
+ var result = execute("SELECT similarity_cosine(value, ?) FROM %s WHERe pk=0", q);
+ assertRows(result, row(1f));
+ result = execute("SELECT similarity_euclidean(value, ?) FROM %s WHERe pk=0", q);
+ assertRows(result, row(1f));
+ execute("SELECT similarity_cosine(?, value) FROM %s WHERE pk=0", q);
+ assertThatThrownBy(() -> execute("SELECT similarity_cosine(?, ?) FROM %s WHERE pk=0", q, q))
+ .hasMessageContaining("Cannot infer type of argument ?");
+
+ // with explicit typing
+ execute("SELECT similarity_cosine((vector<float, 2>) ?, ?) FROM %s WHERE pk=0", q, q);
+ execute("SELECT similarity_cosine(?, (vector<float, 2>) ?) FROM %s WHERE pk=0", q, q);
+ execute("SELECT similarity_cosine((vector<float, 2>) ?, (vector<float, 2>) ?) FROM %s WHERE pk=0", q, q);
+ }
+
+ @Test
+ public void selectSimilarityWithAnn()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ Vector<Float> q = vector(1.5f, 2.5f, 3.5f);
+ var result = execute("SELECT str_val, similarity_cosine(val, ?) FROM %s ORDER BY val ANN OF ? LIMIT 2",
+ q, q);
+
+ assertRowsIgnoringOrder(result,
+ row("A", 0.9987074f),
+ row("B", 0.9993764f));
+ }
+
+ @Test
+ public void castedTerminalFloatVectorFunctions()
+ {
+ createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, value vector<float, 2>)");
+
+ execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(1f, 2f));
+ execute("SELECT similarity_cosine(value, (vector<float, 2>) [1.0, 1.0]) FROM %s WHERE pk=0");
+ execute("SELECT similarity_cosine((vector<float, 2>) [1.0, 1.0], value) FROM %s WHERE pk=0");
+ execute("SELECT similarity_cosine((vector<float, 2>) [1.0, 1.0], (vector<float, 2>) [1.0, 1.0]) FROM %s WHERE pk=0");
+ }
+
+ @Test
+ public void inferredTerminalFloatVectorFunctions() throws Throwable
Review Comment:
There looked to be one group of tests in `inferredTerminalFloatVectorFunctions` that were not covered in `VectorFctsTest`. I have added them to `VectorFctsTest` which has more complete coverage for the different functions and have removed the tests in `VectorTypeTest`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358182918
##########
src/java/org/apache/cassandra/io/util/RandomAccessReader.java:
##########
@@ -71,6 +73,68 @@ private void reBufferAt(long position)
assert buffer.order() == ByteOrder.BIG_ENDIAN : "Buffer must have BIG ENDIAN byte ordering";
}
+ public void readFully(float[] dest) throws IOException
Review Comment:
I realised that when I updated JVector to 1.0.2 and added more code in this class. The original reason I added here was because `bufferHolder` and the constructor were private.
I've changed this and pulled the `RandomAccessReaderAdapter` out into the `vector` package and made it extend `RandomAccessReader`. This looks a lot cleaner and limits the changes here to making `bufferHolder` and the constructor protected.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1333470251
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -664,7 +712,7 @@ public void insertRow(Row row)
@Override
public void updateRow(Row oldRow, Row newRow)
{
- insertRow(newRow);
+ indexContext.getMemtableIndexManager().update(key, oldRow, newRow, memtable);
Review Comment:
This was using `insertRow`, which calls `adjustMemtableSize`. Don't we need that here too?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327569820
##########
src/java/org/apache/cassandra/index/sai/disk/SSTableRowIdPostingList.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.index.sai.disk;
+
+import java.io.IOException;
+
+import com.google.common.base.Preconditions;
+
+import org.apache.cassandra.index.sai.postings.PostingList;
+
+public class SSTableRowIdPostingList implements PostingList
Review Comment:
This entire class seems unused
##########
src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java:
##########
@@ -104,22 +105,20 @@ public void complete(Stopwatch stopwatch) throws IOException
return;
}
- final Iterator<Pair<ByteComparable, LongArrayList>> iterator = rowMapping.merge(memtable);
-
- try (MemtableTermsIterator terms = new MemtableTermsIterator(memtable.getMinTerm(), memtable.getMaxTerm(), iterator))
+ if (indexContext.isVector())
{
- long cellCount = flush(rowMapping.minKey, rowMapping.maxKey, indexContext.getValidator(), terms, rowMapping.maxSSTableRowId);
-
- indexDescriptor.createComponentOnDisk(IndexComponent.COLUMN_COMPLETION_MARKER, indexContext);
-
- indexContext.getIndexMetrics().memtableIndexFlushCount.inc();
-
- long elapsed = stopwatch.elapsed(TimeUnit.MILLISECONDS);
+ flushVectorIndex(rowMapping.minKey, rowMapping.maxKey, start, stopwatch);
Review Comment:
We probably don't need the min/max keys as arguments of `flushVectorIndex`, since it has access to `rowMapping`.
##########
src/java/org/apache/cassandra/index/sai/disk/v1/PerColumnIndexFiles.java:
##########
@@ -71,6 +77,11 @@ public FileHandle balancedTree()
return getFile(IndexComponent.BALANCED_TREE);
}
+ public FileHandle vectors()
+ {
+ return getFile(IndexComponent.VECTOR);
Review Comment:
Maybe the index component should be named `VECTORS`, in the plural, the same as the method?
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/CheckpointingIterator.java:
##########
@@ -0,0 +1,91 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import com.google.common.collect.Iterables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.io.util.FileUtils;
+
+public class CheckpointingIterator extends KeyRangeIterator
+{
+ private static final Logger logger = LoggerFactory.getLogger(CheckpointingIterator.class);
+
+ private final QueryContext context;
+ private final KeyRangeIterator union;
+ private final Iterable<SSTableIndex> referencedIndexes;
+
+ public CheckpointingIterator(KeyRangeIterator wrapped, Iterable<SSTableIndex> referencedIndexes, Iterable<SSTableIndex> referencedAnnIndexesInHybridSearch, QueryContext queryContext)
+ {
+ super(wrapped.getMinimum(), wrapped.getMaximum(), wrapped.getCount());
+
+ this.union = wrapped;
+ if (referencedAnnIndexesInHybridSearch != null)
+ this.referencedIndexes = Iterables.concat(referencedIndexes, referencedAnnIndexesInHybridSearch);
+ else
+ this.referencedIndexes = referencedIndexes;
+ this.context = queryContext;
+ }
+
+ protected PrimaryKey computeNext()
+ {
+ try
+ {
+ return union.hasNext() ? union.next() : endOfData();
+ }
+ finally
+ {
+ context.checkpoint();
+ }
+ }
+
+ protected void performSkipTo(PrimaryKey nextKey)
+ {
+ try
+ {
+ union.skipTo(nextKey);
+ }
+ finally
+ {
+ context.checkpoint();
+ }
+ }
+
+ public void close()
Review Comment:
Nit: add `@Override`
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/CheckpointingIterator.java:
##########
@@ -0,0 +1,91 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import com.google.common.collect.Iterables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.io.util.FileUtils;
+
+public class CheckpointingIterator extends KeyRangeIterator
Review Comment:
Can we add some class JavaDoc?
##########
src/java/org/apache/cassandra/index/sai/disk/v1/SSTableIndexWriter.java:
##########
@@ -82,7 +84,7 @@ public SSTableIndexWriter(IndexDescriptor indexDescriptor, IndexContext indexCon
this.analyzer = indexContext.getAnalyzerFactory().create();
this.limiter = limiter;
this.isIndexValid = isIndexValid;
- this.maxTermSize = indexContext.isFrozen() ? MAX_FROZEN_TERM_SIZE : MAX_STRING_TERM_SIZE;
+ this.maxTermSize = indexContext.isVector() ? MAX_VECTOR_TERM_SIZE : indexContext.isFrozen() ? MAX_FROZEN_TERM_SIZE : MAX_STRING_TERM_SIZE;
Review Comment:
Using parenthesis could ease the reading of the nested ternary operator.
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/CheckpointingIterator.java:
##########
@@ -0,0 +1,91 @@
+/*
+ * 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.index.sai.disk.v1.vector;
+
+import com.google.common.collect.Iterables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.disk.SSTableIndex;
+import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.io.util.FileUtils;
+
+public class CheckpointingIterator extends KeyRangeIterator
+{
+ private static final Logger logger = LoggerFactory.getLogger(CheckpointingIterator.class);
+
+ private final QueryContext context;
+ private final KeyRangeIterator union;
+ private final Iterable<SSTableIndex> referencedIndexes;
+
+ public CheckpointingIterator(KeyRangeIterator wrapped, Iterable<SSTableIndex> referencedIndexes, Iterable<SSTableIndex> referencedAnnIndexesInHybridSearch, QueryContext queryContext)
Review Comment:
Can we break the long line?
```suggestion
public CheckpointingIterator(KeyRangeIterator wrapped,
Iterable<SSTableIndex> referencedIndexes,
Iterable<SSTableIndex> referencedAnnIndexesInHybridSearch,
QueryContext queryContext)
```
##########
src/java/org/apache/cassandra/index/sai/disk/v1/vector/hnsw/BitsUtil.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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.index.sai.disk.v1.vector.hnsw;
+
+import org.apache.lucene.util.Bits;
+import org.cliffc.high_scale_lib.NonBlockingHashMapLong;
+
+import java.util.Set;
+
+public class BitsUtil
+{
+ public static Bits bitsIgnoringDeleted(Bits toAccept, Set<Integer> deletedOrdinals)
+ {
+ return deletedOrdinals.isEmpty()
+ ? toAccept
+ : toAccept == null ? new NoDeletedBits(deletedOrdinals) : new NoDeletedIntersectingBits(toAccept, deletedOrdinals);
+ }
+
+ public static <T> Bits bitsIgnoringDeleted(Bits toAccept, NonBlockingHashMapLong<VectorPostings<T>> postings)
+ {
+ return toAccept == null ? new NoDeletedPostings(postings) : new NoDeletedIntersectingPostings(toAccept, postings);
Review Comment:
```suggestion
return toAccept == null ? new NoDeletedPostings<>(postings) : new NoDeletedIntersectingPostings<>(toAccept, postings);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1334295986
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -135,6 +135,9 @@ private boolean needsReplicaFilteringProtection()
if (command.rowFilter().isEmpty())
return false;
+ if (command.isTopK())
+ return false;
Review Comment:
Thanks for this analysis.
I feel that, for this initial implementation, we should restrict the queries to CL=ONE. Should we allow CL=LOCAL_ONE as well? I'll raise a separate ticket to see how we can improve this.
I'll also add a separate dtest to cover this situation as you have shown above.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1364030760
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,440 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
Review Comment:
Right, and it doesn't make sense to include it in the warning message that is displayed after creating an index. We should instead fail the creation of the index if there are multiple directories, I think.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1348664357
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -281,14 +287,27 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
Selectors selectors = selection.newSelectors(options);
AggregationSpecification aggregationSpec = getAggregationSpec(options);
- ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
- nowInSec, userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+ DataLimits limit = getDataLimits(userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
Review Comment:
I think `PER PARTITION LIMIT` doesn't work with vsearch. That's because the index uses `readCommand.limits().count()` but ignores `readCommand.limits().perPartitionCount()`. This example shows how it returns wrong results:
```java
try (Cluster cluster = init(Cluster.build(2).withConfig(config -> config.with(GOSSIP, NETWORK)).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int, c int, v vector<float, 2>, PRIMARY KEY (k, c))"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex' WITH OPTIONS = { 'similarity_function': 'EUCLIDEAN'}"));
SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
ICoordinator coordinator = cluster.coordinator(1);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (1, 1, [0, 4])"), ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (1, 2, [0, 3])"), ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (2, 1, [0, 2])"), ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (2, 2, [0, 1])"), ALL);
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t PER PARTITION LIMIT 1 LIMIT 100"), ONE),
row(1, 1, vector(0f, 4f)),
row(2, 1, vector(0f, 2f)));
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [0, 0] PER PARTITION LIMIT 1 LIMIT 100"), ONE),
row(2, 2, vector(0f, 1f)),
row(1, 2, vector(0f, 3f))); // returns [[2, 1, [0.0, 2.0]], [1, 1, [0.0, 4.0]]]
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t PER PARTITION LIMIT 1 LIMIT 1"), ONE),
row(1, 1, vector(0f, 4f)));
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [0, 0] PER PARTITION LIMIT 1 LIMIT 1"), ONE),
row(2, 2, vector(0f, 1f)));
}
```
We should probably forbid `PER PARTITION LIMIT` with ANN, and add it to the list of unsupported features in the warning message.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1340216519
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1133,13 +1149,13 @@ public SelectStatement prepare(ClientState state, boolean forView) throws Invali
List<Selectable> selectables = RawSelector.toSelectables(selectClause, table);
boolean containsOnlyStaticColumns = selectOnlyStaticColumns(table, selectables);
- StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, containsOnlyStaticColumns, forView);
+ List<Ordering> orderings = getOrderings(table);
+ StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, orderings, containsOnlyStaticColumns, forView);
// If we order post-query, the sorted column needs to be in the ResultSet for sorting,
// even if we don't ultimately ship them to the client (CASSANDRA-4911).
- Map<ColumnMetadata, Boolean> orderingColumns = getOrderingColumns(table);
- Set<ColumnMetadata> resultSetOrderingColumns = restrictions.keyIsInRelation() ? orderingColumns.keySet()
- : Collections.emptySet();
+ Map<ColumnMetadata, Ordering> orderingColumns = getOrderingColumns(orderings);
+ Set<ColumnMetadata> resultSetOrderingColumns = getResultSetOrdering(restrictions, orderingColumns);
Review Comment:
This failure is reproducible without SAI involved at all. The following test also fails on the `cassandra-5.0` branch:
```java
@Test
public void test2() throws IOException
{
try (Cluster cluster = init(Cluster.build(1).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int, c int, v int, primary key(k, c))"));
cluster.coordinator(1).execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (0, 1, 2)"), ConsistencyLevel.QUORUM);
String query = withKeyspace("SELECT v FROM %s.t WHERE k IN (0, 1) ORDER BY c LIMIT 10");
assertRows(cluster.coordinator(1).execute(query, ConsistencyLevel.ONE), row(2));
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1340232886
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1133,13 +1149,13 @@ public SelectStatement prepare(ClientState state, boolean forView) throws Invali
List<Selectable> selectables = RawSelector.toSelectables(selectClause, table);
boolean containsOnlyStaticColumns = selectOnlyStaticColumns(table, selectables);
- StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, containsOnlyStaticColumns, forView);
+ List<Ordering> orderings = getOrderings(table);
+ StatementRestrictions restrictions = prepareRestrictions(state, table, bindVariables, orderings, containsOnlyStaticColumns, forView);
// If we order post-query, the sorted column needs to be in the ResultSet for sorting,
// even if we don't ultimately ship them to the client (CASSANDRA-4911).
- Map<ColumnMetadata, Boolean> orderingColumns = getOrderingColumns(table);
- Set<ColumnMetadata> resultSetOrderingColumns = restrictions.keyIsInRelation() ? orderingColumns.keySet()
- : Collections.emptySet();
+ Map<ColumnMetadata, Ordering> orderingColumns = getOrderingColumns(orderings);
+ Set<ColumnMetadata> resultSetOrderingColumns = getResultSetOrdering(restrictions, orderingColumns);
Review Comment:
I think both of these make sense because without the ordering column being returned it would be impossible to determine the ordering logic. If the above query is changed to `SELECT v FROM %s.t WHERE k = 0 ORDER BY c LIMIT 10`, then the clustering key is not returned.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341302437
##########
src/java/org/apache/cassandra/service/reads/DataResolver.java:
##########
@@ -135,6 +135,9 @@ private boolean needsReplicaFilteringProtection()
if (command.rowFilter().isEmpty())
return false;
+ if (command.isTopK())
+ return false;
Review Comment:
Great. Now that we don't support queries with CL > ONE the `DataResolver` won't be used in top-k queries. I think we can remove all the changes on this class disabling SRP, RFP and RR if the command is top-k.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341464384
##########
src/java/org/apache/cassandra/cql3/SingleColumnRelation.java:
##########
@@ -264,6 +265,16 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
return new SingleColumnRestriction.LikeRestriction(columnDef, operator, term);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ ColumnMetadata columnDef = table.getExistingColumn(entity);
+ if (!(columnDef.type instanceof VectorType))
Review Comment:
Sorry, I have a whole raft of changes to push. I will get these up later today.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342696896
##########
src/java/org/apache/cassandra/cql3/Operator.java:
##########
@@ -258,6 +258,20 @@ public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteB
{
throw new UnsupportedOperationException();
}
+ },
+ ANN(15)
+ {
+ @Override
+ public String toString()
+ {
+ return "ANN";
+ }
+
+ @Override
+ public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteBuffer rightOperand)
+ {
+ return true;
Review Comment:
Note: If have reverted the change and put an explanation of why it needs to return true here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342754949
##########
src/java/org/apache/cassandra/cql3/Relation.java:
##########
@@ -155,6 +155,8 @@ public final Restriction toRestriction(TableMetadata table, VariableSpecificatio
case LIKE_MATCHES:
case LIKE:
return newLikeRestriction(table, boundNames, relationType);
+// case ANN:
+// return newAnnRestriction(table, boundNames);
Review Comment:
This seems unused. I understand that we shouldn't reach this with ANN. If that's the case, maybe we should add an assertion error here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342557892
##########
src/java/org/apache/cassandra/index/sai/VectorContext.java:
##########
@@ -0,0 +1,168 @@
+/*
+ * 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.index.sai;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnDiskHnsw;
+import org.apache.cassandra.index.sai.disk.v1.vector.hnsw.CassandraOnHeapHnsw;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.lucene.util.Bits;
+
+public class VectorContext
+{
+ public int hnswVectorsAccessed;
+ public int hnswVectorCacheHits;
Review Comment:
These have been removed because they are no longer used by the JVector graph implementation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1374355444
##########
test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java:
##########
@@ -0,0 +1,660 @@
+/*
+ * 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.index.sai.cql;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.junit.Test;
+
+import org.apache.cassandra.config.CassandraRelevantProperties;
+import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.index.sai.StorageAttachedIndex;
+import org.apache.cassandra.service.ClientWarn;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class VectorTypeTest extends VectorTester
+{
+ private static final IPartitioner partitioner = Murmur3Partitioner.instance;
+
+ @Test
+ public void endToEndTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', [5.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (5, 'F', [6.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (6, 'G', [7.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (7, 'H', [8.0, 5.0, 6.0])");
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5");
+ assertThat(result).hasSize(5);
+
+ // some data that only lives in memtable
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (8, 'I', [9.0, 5.0, 6.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (9, 'J', [10.0, 6.0, 7.0])");
+ result = execute("SELECT * FROM %s ORDER BY val ann of [9.5, 5.5, 6.5] LIMIT 5");
+ assertContainsInt(result, "pk", 8);
+ assertContainsInt(result, "pk", 9);
+
+ // data from sstables
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertContainsInt(result, "pk", 1);
+ assertContainsInt(result, "pk", 2);
+ }
+
+ @Test
+ public void warningIsIssuedOnIndexCreation()
+ {
+ ClientWarn.instance.captureWarnings();
+ createTable("CREATE TABLE %s (pk int, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+ List<String> warnings = ClientWarn.instance.getWarnings();
+
+ assertTrue(warnings.size() > 0);
+ assertEquals(StorageAttachedIndex.VECTOR_USAGE_WARNING, warnings.get(0));
+ }
+
+ @Test
+ public void createIndexAfterInsertTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ flush();
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+ }
+
+ public static void assertContainsInt(UntypedResultSet result, String columnName, int columnValue)
+ {
+ for (UntypedResultSet.Row row : result)
+ {
+ if (row.has(columnName))
+ {
+ int value = row.getInt(columnName);
+ if (value == columnValue)
+ {
+ return;
+ }
+ }
+ }
+ throw new AssertionError("Result set does not contain a row with " + columnName + " = " + columnValue);
+ }
+
+ @Test
+ public void testTwoPredicates()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, b, v) VALUES (0, true, [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (1, true, [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (2, false, [3.0, 4.0, 5.0])");
+
+ // the vector given is closest to row 2, but we exclude that row because b=false
+ var result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ // VSTODO assert specific row keys
+ assertThat(result).hasSize(2);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testTwoPredicatesManyRows()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+ for (int i = 0; i < 100; i++)
+ execute("INSERT INTO %s (pk, b, v) VALUES (?, true, ?)",
+ i, vector((float) i, (float) (i + 1), (float) (i + 2)));
+
+ var result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testThreePredicates()
+ {
+ createTable("CREATE TABLE %s (pk int, b boolean, v vector<float, 3>, str text, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(str) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (0, true, [1.0, 2.0, 3.0], 'A')");
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (1, true, [2.0, 3.0, 4.0], 'B')");
+ execute("INSERT INTO %s (pk, b, v, str) VALUES (2, false, [3.0, 4.0, 5.0], 'C')");
+
+ // the vector given is closest to row 2, but we exclude that row because b=false and str!='B'
+ var result = execute("SELECT * FROM %s WHERE b=true AND str='B' ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ // VSTODO assert specific row keys
+ assertThat(result).hasSize(1);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s WHERE b=true AND str='B' ORDER BY v ANN OF [3.1, 4.1, 5.1] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void testSameVectorMultipleRows()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'A', [1.0, 2.0, 3.0])");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+ }
+
+ @Test
+ public void testQueryEmptyTable()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertThat(result).hasSize(0);
+ }
+
+ @Test
+ public void testQueryTableWithNulls()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', null)");
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertThat(result).hasSize(0);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [4.0, 5.0, 6.0])");
+ result = execute("SELECT pk FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 1");
+ assertRows(result, row(1));
+ }
+
+ @Test
+ public void testLimitLessThanInsertedRowCount()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ // Insert more rows than the query limit
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [4.0, 5.0, 6.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [7.0, 8.0, 9.0])");
+
+ // Query with limit less than inserted row count
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void testQueryMoreRowsThanInserted()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+
+ var result = execute("SELECT * FROM %s ORDER BY val ANN OF [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void changingOptionsTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ if (CassandraRelevantProperties.SAI_VECTOR_ALLOW_CUSTOM_PARAMETERS.getBoolean())
+ {
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = " +
+ "{'maximum_node_connections' : 10, 'construction_beam_width' : 200, 'similarity_function' : 'euclidean' }");
+ }
+ else
+ {
+ assertThatThrownBy(() -> createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = " +
+ "{'maximum_node_connections' : 10, 'construction_beam_width' : 200, 'similarity_function' : 'euclidean' }"))
+ .isInstanceOf(InvalidRequestException.class);
+ return;
+ }
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ flush();
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 3");
+ assertThat(result).hasSize(3);
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', [5.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (5, 'F', [6.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (6, 'G', [7.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (7, 'H', [8.0, 5.0, 6.0])");
+
+ flush();
+ compact();
+
+ result = execute("SELECT * FROM %s ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 5");
+ assertThat(result).hasSize(5);
+ }
+
+ @Test
+ public void bindVariablesTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', ?)", vector(2.0f ,3.0f, 4.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', ?)", vector(4.0f, 5.0f, 6.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of ? LIMIT 3", vector(2.5f, 3.5f, 4.5f));
+ assertThat(result).hasSize(3);
+ }
+
+ @Test
+ public void intersectedSearcherTest()
+ {
+ // check that we correctly get back the two rows with str_val=B even when those are not
+ // the closest rows to the query vector
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', ?)", vector(2.0f ,3.0f, 4.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'B', ?)", vector(4.0f, 5.0f, 6.0f));
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', ?)", vector(5.0f, 6.0f, 7.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+
+ flush();
+ result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(2);
+ }
+
+ @Test
+ public void nullVectorTest()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', ?)", vector(1.0f, 2.0f ,3.0f));
+ execute("INSERT INTO %s (pk, str_val) VALUES (1, 'B')"); // no vector
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', ?)", vector(3.0f, 4.0f, 5.0f));
+ execute("INSERT INTO %s (pk, str_val) VALUES (3, 'D')"); // no vector
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'E', ?)", vector(5.0f, 6.0f, 7.0f));
+
+ UntypedResultSet result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(0);
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'A' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+
+ flush();
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'B' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(0);
+
+ result = execute("SELECT * FROM %s WHERE str_val = 'A' ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2");
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void lwtTest()
+ {
+ createTable("CREATE TABLE %s (p int, c int, v text, vec vector<float, 2>, PRIMARY KEY(p, c))");
+ createIndex("CREATE CUSTOM INDEX ON %s(vec) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (p, c, v) VALUES (?, ?, ?)", 0, 0, "test");
+ execute("INSERT INTO %s (p, c, v) VALUES (?, ?, ?)", 0, 1, "00112233445566");
+
+ execute("UPDATE %s SET v='00112233', vec=[0.9, 0.7] WHERE p = 0 AND c = 0 IF v = 'test'");
+
+ UntypedResultSet result = execute("SELECT * FROM %s ORDER BY vec ANN OF [0.1, 0.9] LIMIT 100");
+
+ assertThat(result).hasSize(1);
+ }
+
+ @Test
+ public void twoVectorFieldsTest()
+ {
+ createTable("CREATE TABLE %s (pk int, v2 vector<float, 2>, v3 vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(v2) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(v3) USING 'StorageAttachedIndex'");
+ }
+
+ @Test
+ public void primaryKeySearchTest()
+ {
+ createTable("CREATE TABLE %s (pk int, val vector<float, 3>, i int, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ var N = 5;
+ for (int i = 0; i < N; i++)
+ execute("INSERT INTO %s (pk, val) VALUES (?, ?)", i, vector(1.0f + i, 2.0f + i, 3.0f + i));
+
+ for (int i = 0; i < N; i++)
+ {
+ UntypedResultSet result = execute("SELECT pk FROM %s WHERE pk = ? ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2", i);
+ assertThat(result).hasSize(1);
+ assertRows(result, row(i));
+ }
+
+ flush();
+ for (int i = 0; i < N; i++)
+ {
+ UntypedResultSet result = execute("SELECT pk FROM %s WHERE pk = ? ORDER BY val ann of [2.5, 3.5, 4.5] LIMIT 2", i);
+ assertThat(result).hasSize(1);
+ assertRows(result, row(i));
+ }
+ }
+
+ @Test
+ public void partitionKeySearchTest()
+ {
+ createTable("CREATE TABLE %s (partition int, row int, val vector<float, 2>, PRIMARY KEY(partition, row))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
+
+ var nPartitions = 5;
+ var rowsPerPartition = 10;
+ Map<Integer, List<float[]>> vectorsByPartition = new HashMap<>();
+
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ for (int j = 1; j <= rowsPerPartition; j++)
+ {
+ logger.debug("Inserting partition {} row {}: [{}, {}]", i, j, i, j);
+ execute("INSERT INTO %s (partition, row, val) VALUES (?, ?, ?)", i, j, vector((float) i, (float) j));
+ float[] vector = {(float) i, (float) j};
+ vectorsByPartition.computeIfAbsent(i, k -> new ArrayList<>()).add(vector);
+ }
+ }
+
+ var queryVector = vector(new float[] { 1.5f, 1.5f });
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition, row FROM %s WHERE partition = ? ORDER BY val ann of ? LIMIT 2", i, queryVector);
+ assertThat(result).hasSize(2);
+ assertRowsIgnoringOrder(result,
+ row(i, 1),
+ row(i, 2));
+ }
+
+ flush();
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition, row FROM %s WHERE partition = ? ORDER BY val ann of ? LIMIT 2", i, queryVector);
+ assertThat(result).hasSize(2);
+ assertRowsIgnoringOrder(result,
+ row(i, 1),
+ row(i, 2));
+ }
+ }
+
+ @Test
+ public void clusteringKeyIndexTest()
+ {
+ createTable("CREATE TABLE %s (pk int, ck vector<float, 2>, PRIMARY KEY(pk, ck))");
+ createIndex("CREATE CUSTOM INDEX ON %s(ck) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, ck) VALUES (1, [1.0, 2.0])");
+
+ assertRows(execute("SELECT * FROM %s ORDER BY ck ANN OF [1.0, 2.0] LIMIT 1"), row(1, vector(1.0F, 2.0F)));
+ }
+
+ @Test
+ public void rangeSearchTest() throws Throwable
+ {
+ createTable("CREATE TABLE %s (partition int, val vector<float, 2>, PRIMARY KEY(partition))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function' : 'euclidean'}");
+
+ var nPartitions = 100;
+ Map<Integer, float[]> vectorsByKey = new HashMap<>();
+
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ float[] vector = {(float) i, (float) i};
+ execute("INSERT INTO %s (partition, val) VALUES (?, ?)", i, vector(vector));
+ vectorsByKey.put(i, vector);
+ }
+
+ var queryVector = vector(new float[] { 1.5f, 1.5f });
+ CheckedFunction tester = () -> {
+ for (int i = 1; i <= nPartitions; i++)
+ {
+ UntypedResultSet result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithLowerBound(vectorsByKey.keySet(), i, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithLowerBound(vectorsByKey.keySet(), i, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithUpperBound(vectorsByKey.keySet(), i, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysWithUpperBound(vectorsByKey.keySet(), i, true));
+
+ for (int j = 1; j <= nPartitions; j++)
+ {
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) AND token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, true, j, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) AND token(partition) <= token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, false, j, true));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) >= token(?) AND token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, true, j, false));
+
+ result = execute("SELECT partition FROM %s WHERE token(partition) > token(?) AND token(partition) < token(?) ORDER BY val ann of ? LIMIT 1000", i, j, queryVector);
+ assertThat(keys(result)).containsExactlyInAnyOrderElementsOf(keysInBounds(vectorsByKey.keySet(), i, false, j, false));
+ }
+ }
+ };
+
+ tester.apply();
+
+ flush();
+
+ tester.apply();
+ }
+
+ private Collection<Integer> keys(UntypedResultSet result)
+ {
+ List<Integer> keys = new ArrayList<>(result.size());
+ for (UntypedResultSet.Row row : result)
+ keys.add(row.getInt("partition"));
+ return keys;
+ }
+
+ private Collection<Integer> keysWithLowerBound(Collection<Integer> keys, int leftKey, boolean leftInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getToken(Int32Type.instance.decompose(leftKey)), leftInclusive,
+ partitioner.getMaximumToken().getToken(), true);
+ }
+
+ private Collection<Integer> keysWithUpperBound(Collection<Integer> keys, int rightKey, boolean rightInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getMinimumToken().getToken(), true,
+ partitioner.getToken(Int32Type.instance.decompose(rightKey)), rightInclusive);
+ }
+
+ private Collection<Integer> keysInBounds(Collection<Integer> keys, int leftKey, boolean leftInclusive, int rightKey, boolean rightInclusive)
+ {
+ return keysInTokenRange(keys, partitioner.getToken(Int32Type.instance.decompose(leftKey)), leftInclusive,
+ partitioner.getToken(Int32Type.instance.decompose(rightKey)), rightInclusive);
+ }
+
+ private Collection<Integer> keysInTokenRange(Collection<Integer> keys, Token leftToken, boolean leftInclusive, Token rightToken, boolean rightInclusive)
+ {
+ long left = leftToken.getLongValue();
+ long right = rightToken.getLongValue();
+ return keys.stream()
+ .filter(k -> {
+ long t = partitioner.getToken(Int32Type.instance.decompose(k)).getLongValue();
+ return (left < t || left == t && leftInclusive) && (t < right || t == right && rightInclusive);
+ }).collect(Collectors.toSet());
+ }
+
+ @Test
+ public void selectFloatVectorFunctions()
+ {
+ createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, value vector<float, 2>)");
+
+ // basic functionality
+ Vector<Float> q = vector(1f, 2f);
+ execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector(1f, 2f));
+ execute("SELECT similarity_cosine(value, value) FROM %s WHERE pk=0");
+
+ // type inference checks
+ var result = execute("SELECT similarity_cosine(value, ?) FROM %s WHERe pk=0", q);
+ assertRows(result, row(1f));
+ result = execute("SELECT similarity_euclidean(value, ?) FROM %s WHERe pk=0", q);
+ assertRows(result, row(1f));
+ execute("SELECT similarity_cosine(?, value) FROM %s WHERE pk=0", q);
+ assertThatThrownBy(() -> execute("SELECT similarity_cosine(?, ?) FROM %s WHERE pk=0", q, q))
+ .hasMessageContaining("Cannot infer type of argument ?");
+
+ // with explicit typing
+ execute("SELECT similarity_cosine((vector<float, 2>) ?, ?) FROM %s WHERE pk=0", q, q);
+ execute("SELECT similarity_cosine(?, (vector<float, 2>) ?) FROM %s WHERE pk=0", q, q);
+ execute("SELECT similarity_cosine((vector<float, 2>) ?, (vector<float, 2>) ?) FROM %s WHERE pk=0", q, q);
+ }
+
+ @Test
+ public void selectSimilarityWithAnn()
+ {
+ createTable("CREATE TABLE %s (pk int, str_val text, val vector<float, 3>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'A', [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (2, 'C', [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'D', [4.0, 5.0, 6.0])");
+
+ Vector<Float> q = vector(1.5f, 2.5f, 3.5f);
+ var result = execute("SELECT str_val, similarity_cosine(val, ?) FROM %s ORDER BY val ANN OF ? LIMIT 2",
+ q, q);
+
+ assertRowsIgnoringOrder(result,
+ row("A", 0.9987074f),
+ row("B", 0.9993764f));
+ }
+
+ @Test
+ public void testTwoPredicatesWithUnnecessaryAllowFiltering()
+ {
+ createTable("CREATE TABLE %s (pk int, b int, v vector<float, 3>, PRIMARY KEY(pk, b))");
+ createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(b) USING 'StorageAttachedIndex'");
+
+ execute("INSERT INTO %s (pk, b, v) VALUES (0, 0, [1.0, 2.0, 3.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (1, 2, [2.0, 3.0, 4.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (2, 4, [3.0, 4.0, 5.0])");
+ execute("INSERT INTO %s (pk, b, v) VALUES (3, 6, [4.0, 5.0, 6.0])");
+
+ // Choose a vector closer to b = 0 to ensure that b's restriction is applied.
+ assertRows(execute("SELECT pk FROM %s WHERE b > 2 ORDER BY v ANN OF [1,2,3] LIMIT 2 ALLOW FILTERING;"),
+ row(2), row(3));
+ }
+
+ @Test
+ public void testMultipleVectorsInMemoryWithPredicate()
+ {
+ createTable(KEYSPACE, "CREATE TABLE %s (pk int, val text, vec vector<float, 2>, PRIMARY KEY(pk))");
+ createIndex("CREATE CUSTOM INDEX ON %s(vec) USING 'StorageAttachedIndex'");
+ createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
+
+ // When we search the memtable, we filter out PKs outside the memtable's bounrdaries.
Review Comment:
```suggestion
// When we search the memtable, we filter out PKs outside the memtable's boundaries.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1355023243
##########
src/java/org/apache/cassandra/index/sai/IndexContext.java:
##########
@@ -354,6 +371,11 @@ public boolean isLiteral()
return TypeUtil.isLiteral(getValidator());
}
+ public boolean isVector()
+ {
+ return getValidator().isVector() && ((VectorType<?>)getValidator()).elementType instanceof FloatType;
Review Comment:
There is also an argument that says that it should be treated as a collection and the index type be based on the element type.
For the CEP leaving it as ANN only for float and no other indexes allowed on vectors, and we raise a follow on ticket to support other index types on vectors (probably with a dev discussion first).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370489241
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +71,22 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row ID for a given token or a negative value if not found
+ *
+ * @param token the {@link Token} to lookup
+ * @return the first row ID associated with the {@link Token} or a negative value
+ */
+ long ceiling(Token token);
+
+ /**
+ * Returns the last row ID for a given partition range or the previous row ID if there was no match for the partition range
Review Comment:
Yes, I think the original code uses something `the row ID of the least PK greater than or equal to the given PK` for ceiling, and `the row ID of the greatest PK less than or equal to the given PK` for floor, which sounds clearer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327119333
##########
src/java/org/apache/cassandra/cql3/Ordering.java:
##########
@@ -0,0 +1,187 @@
+/*
+ * 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;
+
+import org.apache.cassandra.cql3.restrictions.SingleColumnRestriction;
+import org.apache.cassandra.cql3.restrictions.SingleRestriction;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.TableMetadata;
+
+/**
+ * A single element of an ORDER BY clause.
+ * <code>ORDER BY ordering1 [, ordering2 [, ...]] </code>
+ * <p>
+ * An ordering comprises an expression that produces the values to compare against each other
+ * and a sorting direction (ASC, DESC).
+ */
+public class Ordering
+{
+ public final Expression expression;
+ public final Direction direction;
+
+ public Ordering(Expression expression, Direction direction)
+ {
+ this.expression = expression;
+ this.direction = direction;
+ }
+
+ public interface Expression
+ {
+ default boolean hasNonClusteredOrdering()
+ {
+ return false;
+ }
+
+ default SingleRestriction toRestriction()
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ ColumnMetadata getColumn();
Review Comment:
All ordering expressions take a column of the constructor and have an identical implementation of `getColumn`. Maybe `Ordering.Expression` could be an abstract class instead of an interface? As in:
```java
public static abstract class Expression
{
protected final ColumnMetadata column;
public Expression(ColumnMetadata column)
{
this.column = column;
}
public boolean hasNonClusteredOrdering()
{
return false;
}
public SingleRestriction toRestriction()
{
throw new UnsupportedOperationException();
}
public ColumnMetadata getColumn()
{
return column;
}
}
```
##########
src/java/org/apache/cassandra/cql3/SingleColumnRelation.java:
##########
@@ -264,6 +265,16 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
return new SingleColumnRestriction.LikeRestriction(columnDef, operator, term);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ ColumnMetadata columnDef = table.getExistingColumn(entity);
+ if (!(columnDef.type instanceof VectorType))
+ throw invalidRequest("ANN is only supported against DENSE FLOAT32 columns");
Review Comment:
The error message should probably mention the CQL vector type, instead of `DENSE FLOAT`. Also, so we have unit tests for this?
##########
src/java/org/apache/cassandra/cql3/SingleColumnRelation.java:
##########
@@ -264,6 +265,16 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
return new SingleColumnRestriction.LikeRestriction(columnDef, operator, term);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ ColumnMetadata columnDef = table.getExistingColumn(entity);
+ if (!(columnDef.type instanceof VectorType))
Review Comment:
Shouldn't we also check that the vector type is for floats?
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1573,6 +1656,39 @@ public int compare(List<ByteBuffer> a, List<ByteBuffer> b)
}
}
+ private static class IndexColumnComparator<T> extends ColumnComparator<List<ByteBuffer>>
+ {
+ private final SingleRestriction restriction;
+ private final int columnIndex;
+
+ // VSTODO maybe cache in prepared statement
+ public IndexColumnComparator(SingleRestriction restriction, int columnIndex)
+ {
+ this.restriction = restriction;
+ this.columnIndex = columnIndex;
+ }
+
+ @Override
+ public boolean indexOrdering()
+ {
+ return true;
+ }
+
+ @Override
+ public Comparator<List<ByteBuffer>> prepareFor(TableMetadata table, QueryOptions options)
+ {
+ Index index = restriction.findSupportingIndex(IndexRegistry.obtain(table));
Review Comment:
This finds the first of those indexes that support the restriction. However, the index selection for the query is made by `SIM#getBestIndexQueryPlanFor`. So I think a query using an index implementation could end up using the comparator provided by another index implementation. This doesn't happen in practice because SAI happens to be the only provided implementation using sorting, albeit it could still happen with custom implementations.
##########
src/java/org/apache/cassandra/cql3/TokenRelation.java:
##########
@@ -118,6 +118,12 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
throw invalidRequest("%s cannot be used with the token function", operator);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ throw invalidRequest("%s cannot be used for toekn relations", operator());
Review Comment:
Typo:s/toekn/token
Also, maybe we should keep the style of the error messages above and use `with the token function` instead of `for token relations`.
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -266,6 +278,8 @@ else if (relation.isLIKE())
// there is restrictions not covered by the PK.
if (!nonPrimaryKeyRestrictions.isEmpty())
{
+ var columnRestrictions = allColumnRestrictions(clusteringColumnsRestrictions, nonPrimaryKeyRestrictions);
Review Comment:
This seems unused
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1327183445
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -276,8 +290,25 @@ else if (relation.isLIKE())
}
if (hasQueriableIndex)
usesSecondaryIndexing = true;
- else if (!allowFiltering && requiresAllowFilteringIfNotSpecified())
- throw invalidRequest(allowFilteringMessage(state));
+ else
+ {
+ if (nonPrimaryKeyRestrictions.hasAnn())
+ {
+ var vectorColumn = nonPrimaryKeyRestrictions.getColumnDefs().stream().filter(c -> c.type.isVector()).findFirst();
Review Comment:
Can't we have multiple vector columns, so the first vector column is not the indexed one? Do we have any tests covering this?
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -276,8 +290,25 @@ else if (relation.isLIKE())
}
if (hasQueriableIndex)
usesSecondaryIndexing = true;
- else if (!allowFiltering && requiresAllowFilteringIfNotSpecified())
- throw invalidRequest(allowFilteringMessage(state));
+ else
Review Comment:
Can we add braces to the `if` block above, so it matches the braced `else` block?
##########
src/java/org/apache/cassandra/db/ReadQuery.java:
##########
@@ -256,4 +256,12 @@ default void maybeValidateIndex()
default void trackWarnings()
{
}
+
+ /**
+ * @return true given read query is a top-k request
Review Comment:
```suggestion
* @return {@code true} if this is a top-k query
```
##########
src/java/org/apache/cassandra/db/ColumnFamilyStore.java:
##########
@@ -1456,6 +1457,12 @@ public void apply(PartitionUpdate update, CassandraWriteContext context, boolean
}
catch (RuntimeException e)
{
+ if (e instanceof InvalidRequestException)
+ {
+ throw new InvalidRequestException(e.getMessage()
Review Comment:
The message construction is identical for `InvalidRequestException` and `RuntimeException`. We could use:
```java
catch (RuntimeException e)
{
String msg = e.getMessage() + " for ks: " + getKeyspaceName() + ", table: " + name;
if (e instanceof InvalidRequestException)
throw new InvalidRequestException(msg, e);
throw new RuntimeException(msg, e);
}
```
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -476,6 +507,38 @@ public boolean usesSecondaryIndexing()
return this.usesSecondaryIndexing;
}
+ /**
+ * This is a hack to push ordering down to indexes.
+ * Indexes are selected based on RowFilter only, so we need to turn orderings into restrictions
+ * so they end up in the row filter.
+ *
+ * @param orderings orderings from the select statement
+ * @return the {@link RestrictionSet} with the added orderings
+ */
+ private RestrictionSet addOrderingRestrictions(List<Ordering> orderings, RestrictionSet restrictionSet)
+ {
+ List<Ordering> annOrderings = orderings.stream().filter(o -> o.expression.hasNonClusteredOrdering()).collect(Collectors.toList());
+
+ if (annOrderings.size() > 1)
+ throw new InvalidRequestException("Cannot specify more than one ANN ordering");
+ else if (annOrderings.size() == 1)
+ {
+ if (orderings.size() > 1)
+ throw new InvalidRequestException("ANN ordering does not support secondary ordering");
Review Comment:
I would say something like `ANN ordering does not support any other ordering`, since ANN could be considered the secondary ordering.
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -1573,6 +1656,39 @@ public int compare(List<ByteBuffer> a, List<ByteBuffer> b)
}
}
+ private static class IndexColumnComparator<T> extends ColumnComparator<List<ByteBuffer>>
Review Comment:
The generic (`<T>`) doesn't seem to be used.
##########
src/java/org/apache/cassandra/db/marshal/VectorType.java:
##########
@@ -365,6 +366,7 @@ public abstract class VectorSerializer extends TypeSerializer<List<T>>
public abstract <V> List<V> split(V buffer, ValueAccessor<V> accessor);
public abstract <V> V serializeRaw(List<V> elements, ValueAccessor<V> accessor);
+ public abstract float[] deserializeFloatArray(ByteBuffer input);
Review Comment:
Does this method duplicate `composeAsFloat`?
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -276,8 +290,25 @@ else if (relation.isLIKE())
}
if (hasQueriableIndex)
usesSecondaryIndexing = true;
- else if (!allowFiltering && requiresAllowFilteringIfNotSpecified())
- throw invalidRequest(allowFilteringMessage(state));
+ else
+ {
+ if (nonPrimaryKeyRestrictions.hasAnn())
+ {
+ var vectorColumn = nonPrimaryKeyRestrictions.getColumnDefs().stream().filter(c -> c.type.isVector()).findFirst();
+ if (vectorColumn.isPresent())
+ {
+ var vc = vectorColumn.get();
+ var hasIndex = indexRegistry.listIndexes().stream().anyMatch(i -> i.dependsOn(vc));
Review Comment:
Can we add an assertion verifying that `indexRegistry` is not null to prevent the warning?
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -822,7 +825,7 @@ private DataLimits getDataLimits(int userLimit,
// If we do post ordering we need to get all the results sorted before we can trim them.
if (aggregationSpec != AggregationSpecification.AGGREGATE_EVERYTHING)
{
- if (!needsPostQueryOrdering())
+ if (!needsToSkipUserLimit())
Review Comment:
I think it's the only usage of `needsToSkipUserLimit`, so it can probably be inlined here.
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -288,7 +291,7 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
query.trackWarnings();
ResultMessage.Rows rows;
- if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize)))
+ if (aggregationSpec == null && (pageSize <= 0 || (query.limits().count() <= pageSize) || query.isTopK()))
Review Comment:
`StorageAttachedIndex#validate` makes sure that we cannot use ANN without a `LIMIT`. However, paging is disabled in the CQL layer, so any index implementation could provide support for ANN. I think that the existence of a `LIMIT` should be checked here in `SelectStatement`.
As for the specific limit, 100 by default, I guess it could also be checked on the generic CQL layer because it affects the size of the pages. Should it also be checked by the specific index implementation?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] adelapena commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1341318727
##########
test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java:
##########
@@ -0,0 +1,475 @@
+/*
+ * 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.distributed.test.sai;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.distributed.api.ConsistencyLevel;
+import org.apache.cassandra.distributed.test.TestBaseImpl;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
+
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class VectorDistributedTest extends TestBaseImpl
+{
+
+ @Rule
+ public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher();
+
+ private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}";
+ private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)";
+ private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'";
+
+ private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION;
+
+ private static final String INVALID_LIMIT_MESSAGE = "Use of ANN OF in an ORDER BY clause requires a LIMIT that is not greater than 1000";
+
+ private static final double MIN_RECALL = 0.8;
+
+ private static final int NUM_REPLICAS = 3;
+ private static final int RF = 2;
+
+ private static final AtomicInteger seq = new AtomicInteger();
+ private static String table;
+
+ private static Cluster cluster;
+
+ private static int dimensionCount;
+
+ @BeforeClass
+ public static void setupCluster() throws Exception
+ {
+ cluster = Cluster.build(NUM_REPLICAS)
+ .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens
+ .withDataDirCount(1) // VSTODO vector memtable flush doesn't support multiple directories yet
+ .withConfig(config -> config.with(GOSSIP).with(NETWORK))
+ .start();
+
+ cluster.schemaChange(withKeyspace(String.format(CREATE_KEYSPACE, RF)));
+ }
+
+ @AfterClass
+ public static void closeCluster()
+ {
+ if (cluster != null)
+ cluster.close();
+ }
+
+ @Before
+ public void before()
+ {
+ table = "table_" + seq.getAndIncrement();
+ dimensionCount = SAITester.getRandom().nextIntBetween(100, 2048);
+ }
+
+ @After
+ public void after()
+ {
+ cluster.schemaChange(formatQuery("DROP TABLE IF EXISTS %s"));
+ }
+
+ @Test
+ public void topKQueryFailsUnlessAtConsistencyLevelOne()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ float[] queryVector = randomVector();
+ assertThatThrownBy(() -> execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit, ConsistencyLevel.QUORUM))
+ .hasMessage(SelectStatement.TOPK_CONSISTENCY_LEVEL_ERROR);
+
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit, ConsistencyLevel.ONE);
+ assertThat(result).hasNumberOfRows(limit);
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecall = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecall = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithoutLimit(randomVector(), vectorCount))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ int pageSize = SAITester.getRandom().nextIntBetween(40, 70);
+ limit = SAITester.getRandom().nextIntBetween(20, 50);
+ result = searchWithPageAndLimit(queryVector, pageSize, limit);
+
+ resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double memtableRecallWithPaging = getRecall(vectors, queryVector, resultVectors);
+ assertThat(memtableRecallWithPaging).isGreaterThanOrEqualTo(MIN_RECALL);
+
+ assertThatThrownBy(() -> searchWithPageWithoutLimit(randomVector()))
+ .hasMessageContaining(INVALID_LIMIT_MESSAGE);
+
+ // query on-disk index
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ limit = Math.min(SAITester.getRandom().nextIntBetween(10, 50), vectors.size());
+ queryVector = randomVector();
+ result = searchWithLimit(queryVector, limit);
+ double sstableRecall = getRecall(vectors, queryVector, getVectors(result));
+ assertThat(sstableRecall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testMultiSSTablesVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+ // disable compaction
+ String tableName = table;
+ cluster.forEach(n -> n.runOnInstance(() -> {
+ Keyspace keyspace = Keyspace.open(KEYSPACE);
+ keyspace.getColumnFamilyStore(tableName).disableAutoCompaction();
+ }));
+
+ int vectorCountPerSSTable = SAITester.getRandom().nextIntBetween(200, 500);
+ int sstableCount = SAITester.getRandom().nextIntBetween(3, 5);
+ List<float[]> allVectors = new ArrayList<>(sstableCount * vectorCountPerSSTable);
+
+ int pk = 0;
+ for (int i = 0; i < sstableCount; i++)
+ {
+ List<float[]> vectors = generateVectors(vectorCountPerSSTable);
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ allVectors.addAll(vectors);
+ cluster.forEach(n -> n.flush(KEYSPACE));
+ }
+
+ // query multiple sstable indexes in multiple node
+ int limit = Math.min(SAITester.getRandom().nextIntBetween(50, 100), allVectors.size());
+ float[] queryVector = randomVector();
+ Object[][] result = searchWithLimit(queryVector, limit);
+
+ // expect recall to be at least 0.8
+ List<float[]> resultVectors = getVectors(result);
+ assertDescendingScore(queryVector, resultVectors);
+ double recall = getRecall(allVectors, queryVector, getVectors(result));
+ assertThat(recall).isGreaterThanOrEqualTo(MIN_RECALL);
+ }
+
+ @Test
+ public void testPartitionRestrictedVectorSearch()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = generateVectors(vectorCount);
+
+ int pk = 0;
+ for (float[] vector : vectors)
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ", " + vectorString(vector) + " )");
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, vectors);
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key = SAITester.getRandom().nextIntBetween(0, vectorCount - 1);
+ float[] queryVector = randomVector();
+ searchByKeyWithLimit(key, queryVector, vectors);
+ }
+ }
+
+ @Test
+ public void rangeRestrictedTest()
+ {
+ cluster.schemaChange(formatQuery(String.format(CREATE_TABLE, dimensionCount)));
+ cluster.schemaChange(formatQuery(String.format(CREATE_INDEX, "val")));
+ SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
+
+ int vectorCount = SAITester.getRandom().nextIntBetween(500, 1000);
+ List<float[]> vectors = IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+
+ int pk = 0;
+ Multimap<Long, float[]> vectorsByToken = ArrayListMultimap.create();
+ for (float[] vector : vectors)
+ {
+ vectorsByToken.put(Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(pk)).getLongValue(), vector);
+ execute("INSERT INTO %s (pk, val) VALUES (" + (pk++) + ',' + vectorString(vector) + " )");
+ }
+
+ // query memtable index
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+
+ cluster.forEach(n -> n.flush(KEYSPACE));
+
+ // query on-disk index with existing key:
+ for (int executionCount = 0; executionCount < 50; executionCount++)
+ {
+ int key1 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token1 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key1)).getLongValue();
+ int key2 = SAITester.getRandom().nextIntBetween(1, vectorCount * 2);
+ long token2 = Murmur3Partitioner.instance.getToken(Int32Type.instance.decompose(key2)).getLongValue();
+
+ long minToken = Math.min(token1, token2);
+ long maxToken = Math.max(token1, token2);
+ List<float[]> expected = vectorsByToken.entries().stream()
+ .filter(e -> e.getKey() >= minToken && e.getKey() <= maxToken)
+ .map(Map.Entry::getValue)
+ .collect(Collectors.toList());
+
+ float[] queryVector = randomVector();
+
+ List<float[]> resultVectors = searchWithRange(queryVector, minToken, maxToken, expected.size());
+ if (expected.isEmpty())
+ assertThat(resultVectors).isEmpty();
+ else
+ {
+ double recall = getRecall(resultVectors, queryVector, expected);
+ assertThat(recall).isGreaterThanOrEqualTo(0.8);
+ }
+ }
+ }
+
+ private List<float[]> searchWithRange(float[] queryVector, long minToken, long maxToken, int expectedSize)
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE token(pk) <= " + maxToken + " AND token(pk) >= " + minToken + " ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT 1000");
+ assertThat(result).hasNumberOfRows(expectedSize);
+ return getVectors(result);
+ }
+
+ private Object[][] searchWithLimit(float[] queryVector, int limit)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit);
+ assertThat(result).hasNumberOfRows(limit);
+ return result;
+ }
+
+ private void searchWithoutLimit(float[] queryVector, int results)
+ {
+ Object[][] result = execute("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector));
+ assertThat(result).hasNumberOfRows(results);
+ }
+
+
+ private void searchWithPageWithoutLimit(float[] queryVector)
+ {
+ executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector), 10);
+ }
+
+ private Object[][] searchWithPageAndLimit(float[] queryVector, int pageSize, int limit)
+ {
+ // we don't know how many will be returned in case of paging, because coordinator resumes from last-returned-row's partiton
+ return executeWithPaging("SELECT val FROM %s ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit, pageSize);
+ }
+
+ private void searchByKeyWithLimit(int key, float[] queryVector, List<float[]> vectors)
+ {
+ Object[][] result = execute("SELECT val FROM %s WHERE pk = " + key + " ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT 1");
+ assertThat(result).hasNumberOfRows(1);
+ float[] output = getVectors(result).get(0);
+ assertThat(output).isEqualTo(vectors.get(key));
+ }
+
+ private void assertDescendingScore(float[] queryVector, List<float[]> resultVectors)
+ {
+ float prevScore = -1;
+ for (float[] current : resultVectors)
+ {
+ float score = function.compare(current, queryVector);
+ if (prevScore >= 0)
+ assertThat(score).isLessThanOrEqualTo(prevScore);
+
+ prevScore = score;
+ }
+ }
+
+ private double getRecall(List<float[]> vectors, float[] query, List<float[]> result)
+ {
+ List<float[]> sortedVectors = new ArrayList<>(vectors);
+ sortedVectors.sort((a, b) -> Double.compare(function.compare(b, query), function.compare(a, query)));
+
+ assertThat(sortedVectors).containsAll(result);
+
+ List<float[]> nearestNeighbors = sortedVectors.subList(0, result.size());
+
+ int matches = 0;
+ for (float[] in : nearestNeighbors)
+ {
+ for (float[] out : result)
+ {
+ if (Arrays.compare(in, out) ==0)
+ {
+ matches++;
+ break;
+ }
+ }
+ }
+
+ return matches * 1.0 / result.size();
+ }
+
+ private List<float[]> generateVectors(int vectorCount)
+ {
+ return IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+ }
+
+ private List<float[]> getVectors(Object[][] result)
Review Comment:
Nit: add `@SuppressWarnings("unchecked")`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342744090
##########
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java:
##########
@@ -426,6 +426,7 @@ public enum CassandraRelevantProperties
SAI_LATEST_VERSION("cassandra.sai.latest_version", "aa"),
SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size_kb", "5"),
SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size_kb", "1"),
+ SAI_MAX_VECTOR_TERM_SIZE("cassandra.sai.max_vector_term_size_kb", "16"),
Review Comment:
I think this property and the two other above (`MAX_STRING_TERM_SIZE` and `MAX_FROZEN_TERM_SIZE`) should use the new config standard where the config the properties don't need to use a fixed unit. Instead, the size unit can be expressed in the user-specified value:
```java
SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size", "5KiB"),
SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size", "1KiB"),
SAI_MAX_VECTOR_TERM_SIZE("cassandra.sai.max_vector_term_size", "16KiB"),
```
Then callers can use `SAI_MAX_VECTOR_TERM_SIZE.getSizeInBytes()` to get the value in bytes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1342862641
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -109,6 +112,9 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(SelectStatement.logger, 1, TimeUnit.MINUTES);
public static final int DEFAULT_PAGE_SIZE = 10000;
+ public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top-k queries can only be run with consistency level one. Consistency level %s was used.";
Review Comment:
I think we are missing the client warning saying that vector indexes are experimental. That warning message should probably be quite similar to`SASIIndex.USAGE_WARNING` and `View.USAGE_WARNING`. We could also include information about limitations. For example:
```
SAI ANN indexes on vector columns are experimental and are not recommended for production use.
They cannot be queried with a read consistency level higher than ONE/LOCAL_ONE.
They cannot use paging either, nor be used in aggregation queries.
These limitations are planned to be removed in future versions.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1348664357
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -281,14 +287,27 @@ public ResultMessage.Rows execute(QueryState state, QueryOptions options, long q
Selectors selectors = selection.newSelectors(options);
AggregationSpecification aggregationSpec = getAggregationSpec(options);
- ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
- nowInSec, userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
+ DataLimits limit = getDataLimits(userLimit, userPerPartitionLimit, pageSize, aggregationSpec);
Review Comment:
I think `PER PARTITION LIMIT` doesn't work with vsearch. That's because index uses `readCommand.limits().count()` but ignores `readCommand.limits().perPartitionCount()`. This example shows how it returns wrong results:
```java
try (Cluster cluster = init(Cluster.build(2).withConfig(config -> config.with(GOSSIP, NETWORK)).start()))
{
cluster.schemaChange(withKeyspace("CREATE TABLE %s.t (k int, c int, v vector<float, 2>, PRIMARY KEY (k, c))"));
cluster.schemaChange(withKeyspace("CREATE CUSTOM INDEX ON %s.t(v) USING 'StorageAttachedIndex' WITH OPTIONS = { 'similarity_function': 'EUCLIDEAN'}"));
SAIUtil.waitForIndexQueryable(cluster, KEYSPACE);
ICoordinator coordinator = cluster.coordinator(1);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (1, 1, [0, 4])"), ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (1, 2, [0, 3])"), ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (2, 1, [0, 2])"), ALL);
coordinator.execute(withKeyspace("INSERT INTO %s.t (k, c, v) VALUES (2, 2, [0, 1])"), ALL);
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t PER PARTITION LIMIT 1 LIMIT 100"), ONE),
row(1, 1, vector(0f, 4f)),
row(2, 1, vector(0f, 2f)));
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [0, 0] PER PARTITION LIMIT 1 LIMIT 100"), ONE),
row(2, 2, vector(0f, 1f)),
row(1, 2, vector(0f, 3f))); // returns [[2, 1, [0.0, 2.0]], [1, 1, [0.0, 4.0]]]
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t PER PARTITION LIMIT 1 LIMIT 1"), ONE),
row(1, 1, vector(0f, 4f)));
assertRows(coordinator.execute(withKeyspace("SELECT * FROM %s.t ORDER BY v ANN OF [0, 0] PER PARTITION LIMIT 1 LIMIT 1"), ONE),
row(2, 2, vector(0f, 1f)));
}
```
We should probably forbid it and add it to the list of unsupported features in the warning message.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1335937452
##########
src/java/org/apache/cassandra/cql3/Ordering.java:
##########
@@ -0,0 +1,187 @@
+/*
+ * 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;
+
+import org.apache.cassandra.cql3.restrictions.SingleColumnRestriction;
+import org.apache.cassandra.cql3.restrictions.SingleRestriction;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.TableMetadata;
+
+/**
+ * A single element of an ORDER BY clause.
+ * <code>ORDER BY ordering1 [, ordering2 [, ...]] </code>
+ * <p>
+ * An ordering comprises an expression that produces the values to compare against each other
+ * and a sorting direction (ASC, DESC).
+ */
+public class Ordering
+{
+ public final Expression expression;
+ public final Direction direction;
+
+ public Ordering(Expression expression, Direction direction)
+ {
+ this.expression = expression;
+ this.direction = direction;
+ }
+
+ public interface Expression
+ {
+ default boolean hasNonClusteredOrdering()
+ {
+ return false;
+ }
+
+ default SingleRestriction toRestriction()
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ ColumnMetadata getColumn();
+ }
+
+ /**
+ * Represents a single column in
+ * <code>ORDER BY column</code>
+ */
+ public static class SingleColumn implements Expression
+ {
+ public final ColumnMetadata column;
+
+ public SingleColumn(ColumnMetadata column)
+ {
+ this.column = column;
+ }
+
+ @Override
+ public ColumnMetadata getColumn()
+ {
+ return column;
+ }
+ }
+
+ /**
+ * An expression used in Approximate Nearest Neighbor ordering.
+ * <code>ORDER BY column ANN OF value</code>
+ */
+ public static class Ann implements Expression
+ {
+ final ColumnMetadata column;
+ final Term vectorValue;
+
+ public Ann(ColumnMetadata column, Term vectorValue)
+ {
+ this.column = column;
+ this.vectorValue = vectorValue;
+ }
+
+ @Override
+ public boolean hasNonClusteredOrdering()
+ {
+ return true;
+ }
+
+ @Override
+ public SingleRestriction toRestriction()
+ {
+ return new SingleColumnRestriction.AnnRestriction(column, vectorValue);
+ }
+
+ @Override
+ public ColumnMetadata getColumn()
+ {
+ return column;
+ }
+ }
+
+ public enum Direction
+ {ASC, DESC}
+
+
+ /**
+ * Represents the AST of a single element in the ORDER BY clause.
Review Comment:
It stands for abstract syntax tree, and I've changed it to that. I realise it's not obvious but it is a term that is used in relation to antlr, maybe not by us but in their documentation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1336149186
##########
src/java/org/apache/cassandra/cql3/TokenRelation.java:
##########
@@ -118,6 +118,12 @@ protected Restriction newLikeRestriction(TableMetadata table, VariableSpecificat
throw invalidRequest("%s cannot be used with the token function", operator);
}
+ @Override
+ protected Restriction newAnnRestriction(TableMetadata table, VariableSpecifications boundNames)
+ {
+ throw invalidRequest("%s cannot be used for toekn relations", operator());
Review Comment:
As previous message, this goes away.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
[GitHub] [cassandra] mike-tr-adamson commented on a diff in pull request #2673: CASSANDRA-18715 Add support for a vector search index in SAI
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1337139772
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -483,7 +486,7 @@ else if (restrictions.keyIsInRelation())
// Please note that the isExhausted state of the pager only gets updated when we've closed the page, so this
// shouldn't be moved inside the 'try' above.
- if (!pager.isExhausted())
+ if (!pager.isExhausted() && !pager.pager.isTopK())
Review Comment:
The query issues were a 2 part problem in the `VectorQueryController`.
1. The code to establish whether a query was a hybrid query or not did not consider whether the other expressions in the query were indexed or not.
2. The code to build the iterators was not checking for indexed expressions correctly.
The `UOE` on flush is not reproducing locally, and the test with flush is running correctly. I can only assume that this has been fixed elsewhere.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1360579083
##########
src/java/org/apache/cassandra/cql3/statements/SelectStatement.java:
##########
@@ -109,6 +112,15 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(SelectStatement.logger, 1, TimeUnit.MINUTES);
public static final int DEFAULT_PAGE_SIZE = 10000;
+ public static final String TOPK_CONSISTENCY_LEVEL_ERROR = "Top-K queries can only be run with consistency level ONE/LOCAL_ONE. Consistency level %s was used.";
+ public static final String TOPK_LIMIT_ERROR = "Top-K queries must have a limit specified and the limit must be less than the query page size";
+ public static final String TOPK_PARTITION_LIMIT_ERROR = "Top-K queries do not support per-partition limits";
+ public static final String TOPK_AGGREGATION_ERROR = "Top-K queries can not be run with aggregation";
+ public static final String TOPK_CONSISTENCY_LEVEL_WARNING = "Top-K queries can only be run with consistency level ONE " +
+ "/ LOCAL_ONE / NODE_LOCAL. Consistency level %s was used. " +
+ "Downgrading to consistency level %s.";
+ public static final String TOPK_PAGE_SIZE_WARNING = "Top-K queries do not support paging and the page size is set to %d, " +
+ "which is less than LIMIT %d. The page size has been set to %d to match the LIMIT.";
Review Comment:
Nit: the two last arguments are identical, so we can use "%<d" for the last one:
```suggestion
public static final String TOPK_PAGE_SIZE_WARNING = "Top-K queries do not support paging and the page size is set to %d, " +
"which is less than LIMIT %d. The page size has been set to %<d to match the LIMIT.";
```
This will make `String.format` calls only need two args:
```java
ClientWarn.instance.warn(String.format(TOPK_PAGE_SIZE_WARNING, oldPageSize, limit.count()));
```
```java
assertEquals(String.format(SelectStatement.TOPK_PAGE_SIZE_WARNING, 9, 10), ClientWarn.instance.getWarnings().get(0));
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1356807145
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,29 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
Review Comment:
Agreed, and I was thinking about the "most of" as well because there are some of these limitations that are unlikely to be lifted.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358126158
##########
src/java/org/apache/cassandra/cql3/restrictions/SingleColumnRestriction.java:
##########
@@ -804,4 +824,70 @@ else if (ByteBufferUtil.startsWith(value, LIKE_WILDCARD))
return Pair.create(operator, newValue);
}
}
+ public static final class AnnRestriction extends SingleColumnRestriction
+ {
+ private final Term value;
+
+ public AnnRestriction(ColumnMetadata columnDef, Term value)
+ {
+ super(columnDef);
+ this.value = value;
+ }
+
+ public ByteBuffer value(QueryOptions options)
+ {
+ return value.bindAndGet(options);
+ }
+
+ @Override
+ public boolean isANN()
+ {
+ return true;
+ }
+
+ @Override
+ public void addFunctionsTo(List<Function> functions)
+ {
+ value.addFunctionsTo(functions);
+ }
+
+ @Override
+ MultiColumnRestriction toMultiColumnRestriction()
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void addToRowFilter(RowFilter filter,
+ IndexRegistry indexRegistry,
+ QueryOptions options)
+ {
+ filter.add(columnDef, Operator.ANN, value.bindAndGet(options));
+ }
+
+ @Override
+ public MultiCBuilder appendTo(MultiCBuilder builder, QueryOptions options)
+ {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public String toString()
+ {
+ return String.format("ANN(%s)", value);
+ }
+
+ @Override
+ public SingleRestriction doMergeWith(SingleRestriction otherRestriction)
+ {
+ // VSTODO not sure if this is right
Review Comment:
It was obsolete but was throwing the wrong exception. I have changed the implementation to:
```
throw invalidRequest("%s cannot be restricted by more than one relation in an ANN ordering", columnDef.name);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1358122126
##########
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java:
##########
@@ -424,8 +424,9 @@ public enum CassandraRelevantProperties
SAI_INTERSECTION_CLAUSE_LIMIT("cassandra.sai.intersection_clause_limit", "2"),
/** Latest version to be used for SAI index writing */
SAI_LATEST_VERSION("cassandra.sai.latest_version", "aa"),
- SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size_kb", "5"),
- SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size_kb", "1"),
+ SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size", "5KiB"),
+ SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size", "1KiB"),
+ SAI_MAX_VECTOR_TERM_SIZE("cassandra.sai.max_vector_term_size", "16KiB"),
Review Comment:
I have set it to 32KiB for now, but there is the potential (later) to remove the term size restriction entirely. This is only in place because the in-memory index builders had a size limitation. The numeric index now uses a trie for the in-memory, and we have a ticket to move the literal index to use the same. This would mean we could remove these and replace them with guardrails instead.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357164060
##########
src/java/org/apache/cassandra/io/util/RandomAccessReader.java:
##########
@@ -71,6 +73,68 @@ private void reBufferAt(long position)
assert buffer.order() == ByteOrder.BIG_ENDIAN : "Buffer must have BIG ENDIAN byte ordering";
}
+ public void readFully(float[] dest) throws IOException
Review Comment:
we have a wrapper class for RAR elsewhere in this patch, but here we're implementing most of it directly, do we need the wrapper?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "jbellis (via GitHub)" <gi...@apache.org>.
jbellis commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1357029954
##########
src/java/org/apache/cassandra/cql3/restrictions/StatementRestrictions.java:
##########
@@ -63,6 +69,14 @@ public final class StatementRestrictions
"Executing this query despite the performance unpredictability with ALLOW FILTERING has been disabled " +
"by the allow_filtering_enabled property in cassandra.yaml";
+ public static final String ANN_REQUIRES_INDEX_MESSAGE = "ANN ordering by vector requires the column to be indexed";
+
+ public static final String VECTOR_INDEXES_ANN_ONLY_MESSAGE = "Vector indexes only support ANN queries";
+
+ public static final String ANN_ONLY_SUPPORTED_ON_VECTOR_MESSAGE = "ANN ordering is only supported on float vector indexes";
+
+ public static final String ANN_REQUIRES_INDEXED_FILTERING_MESSAGE = "ANN ordered queries can only be filtered with indexed columns";
Review Comment:
Might want to grab the improvements from ds vsearch dc3cc020
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370119154
##########
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeIterator.java:
##########
@@ -43,6 +43,11 @@ protected KeyRangeIterator(Builder.Statistics statistics)
this(statistics.min, statistics.max, statistics.count);
}
+ public KeyRangeIterator(KeyRangeIterator range)
+ {
+ this(range == null ? null : range.min, range == null ? null : range.max, range == null ? -1 : range.count);
Review Comment:
Nit: can we break the args for better readability?
```suggestion
this(range == null ? null : range.min,
range == null ? null : range.max,
range == null ? -1 : range.count);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370278525
##########
src/java/org/apache/cassandra/index/sai/utils/AtomicRatio.java:
##########
@@ -0,0 +1,57 @@
+/*
+ * 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.index.sai.utils;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class AtomicRatio
Review Comment:
This could use some class JavaDoc
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1371838395
##########
src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java:
##########
@@ -90,17 +99,35 @@
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.IndexMetadata;
import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.OpOrder;
+import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K;
+
public class StorageAttachedIndex implements Index
{
public static final String NAME = "sai";
+
+ public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" +
Review Comment:
Right. There shouldn't really have been any docs committed yet until the CEP is merged. I'll get a ticket raised when we merge.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "mike-tr-adamson (via GitHub)" <gi...@apache.org>.
mike-tr-adamson commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1371668822
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +71,22 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row ID for a given token or a negative value if not found
+ *
+ * @param token the {@link Token} to lookup
+ * @return the first row ID associated with the {@link Token} or a negative value
+ */
+ long ceiling(Token token);
+
+ /**
+ * Returns the last row ID for a given partition range or the previous row ID if there was no match for the partition range
Review Comment:
I have changed the javadoc for these as follows:
- ceiling - Returns the first row ID of the nearest {@link Token} greater than or equal to the given {@link Token}, or a negative value if not found
- floor - Returns the last row ID of the nearest {@link Token} less than or equal to the given {@link Token}, or a negative value if the {@link Token} is at its minimum value
I believe this is the closest to reality.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370102760
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +71,22 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row ID for a given token or a negative value if not found
+ *
+ * @param token the {@link Token} to lookup
+ * @return the first row ID associated with the {@link Token} or a negative value
+ */
+ long ceiling(Token token);
+
+ /**
+ * Returns the last row ID for a given partition range or the previous row ID if there was no match for the partition range
Review Comment:
It would be "for a given token" instead of "for a given partition range". Also, aren't the description and the functions inverted, so floor should return the first ID and ceiling the last one?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org
Re: [PR] CASSANDRA-18715 Add support for a vector search index in SAI [cassandra]
Posted by "adelapena (via GitHub)" <gi...@apache.org>.
adelapena commented on code in PR #2673:
URL: https://github.com/apache/cassandra/pull/2673#discussion_r1373038899
##########
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java:
##########
@@ -70,6 +71,22 @@ default void close()
*/
long rowIdFromPrimaryKey(PrimaryKey key);
+ /**
+ * Returns the first row ID for a given token or a negative value if not found
+ *
+ * @param token the {@link Token} to lookup
+ * @return the first row ID associated with the {@link Token} or a negative value
+ */
+ long ceiling(Token token);
+
+ /**
+ * Returns the last row ID for a given partition range or the previous row ID if there was no match for the partition range
Review Comment:
It looks great now, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org