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

Refactor SAI ANN query execution to use score ordered iterators for correctness and speed 

Rewrites ANN search query execution logic to more efficiently merge graph search results using
similarity score-ordered (descending) iterators to merge segments efficiently. Allows for reduced memory
consumption during queries, reduced impact of overwrites and tombstones, selective re-querying of
minimally necessary graphs, and reduced shuffling of PrimaryKey objects.

patch by Michael Marshall; reviewed by Caleb Rackliffe and Michael Semb Wever for CASSANDRA-20086
This commit is contained in:
Michael Marshall 2026-03-17 12:42:52 -05:00 committed by Caleb Rackliffe
parent a2e55684de
commit 23ec1c8a3f
71 changed files with 3681 additions and 1173 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGES.txt
View File

@ -1,4 +1,5 @@
5.0.7
* Refactor SAI ANN query execution to use score ordered iterators for correctness and speed (CASSANDRA-20086)
* Disallow binding an identity to a superuser when the user is a regular user (CASSANDRA-21219)
* Fix ConcurrentModificationException in compaction garbagecollect (CASSANDRA-21065)
* Dynamically skip sharding L0 when SAI Vector index present (CASSANDRA-19661)

17
src/java/org/apache/cassandra/config/CassandraRelevantProperties.java
View File

@ -460,15 +460,20 @@ public enum CassandraRelevantProperties
/** Whether to allow the user to specify custom options to the hnsw index */
SAI_VECTOR_ALLOW_CUSTOM_PARAMETERS("cassandra.sai.vector.allow_custom_parameters", "false"),
/**
* The maximum number of primary keys that a WHERE clause may materialize before the query planner switches
* from a search-then-sort execution strategy to an order-by-then-filter strategy. Increasing this limit allows
* more primary keys to be buffered in memory, enabling either (a) brute-force sorting or (b) graph traversal
* with a restrictive filter that admits only nodes whose primary keys matched the WHERE clause.
*
* Note also that the SAI_INTERSECTION_CLAUSE_LIMIT is applied to the WHERE clause before using a search to
* build a potential result set for search-then-sort query execution.
*/
SAI_VECTOR_SEARCH_MAX_MATERIALIZE_KEYS("cassandra.sai.vector_search.max_materialized_keys", "16000"),
/** Controls the maximum top-k limit for vector search */
SAI_VECTOR_SEARCH_MAX_TOP_K("cassandra.sai.vector_search.max_top_k", "1000"),
/**
* Controls the maximum number of PrimaryKeys that will be read into memory at one time when ordering/limiting
* the results of an ANN query constrained by non-ANN predicates.
*/
SAI_VECTOR_SEARCH_ORDER_CHUNK_SIZE("cassandra.sai.vector_search.order_chunk_size", "100000"),
SCHEMA_PULL_INTERVAL_MS("cassandra.schema_pull_interval_ms", "60000"),
SCHEMA_UPDATE_HANDLER_FACTORY_CLASS("cassandra.schema.update_handler_factory.class"),
SEARCH_CONCURRENCY_FACTOR("cassandra.search_concurrency_factor", "1"),

37
src/java/org/apache/cassandra/db/CellSourceIdentifier.java Normal file
View File

@ -0,0 +1,37 @@
/*
* 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.db;
/**
* A cell's source data object. Can be used to determine if two cells originated from the same object, e.g. memtable
* or sstable.
*/
public interface CellSourceIdentifier
{
/**
* Returns true iff this and other CellSourceIdentifier are equal, indicating that the cell are from the same
* source.
* @param other the other source with which to compare
* @return true if the two sources are equal
*/
default boolean isEqualSource(CellSourceIdentifier other)
{
return this.equals(other);
}
}

55
src/java/org/apache/cassandra/db/SinglePartitionReadCommand.java
View File

@ -26,6 +26,7 @@ import java.util.List;
import java.util.NavigableSet;
import java.util.TreeSet;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;
import com.google.common.annotations.VisibleForTesting;
@ -51,6 +52,7 @@ import org.apache.cassandra.db.partitions.PartitionIterator;
import org.apache.cassandra.db.partitions.PartitionIterators;
import org.apache.cassandra.db.partitions.SingletonUnfilteredPartitionIterator;
import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
import org.apache.cassandra.db.rows.BaseRowIterator;
import org.apache.cassandra.db.rows.Cell;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.db.rows.Rows;
@ -658,10 +660,26 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
assert executionController != null && executionController.validForReadOn(cfs);
Tracing.trace("Executing single-partition query on {}", cfs.name);
return queryMemtableAndDiskInternal(cfs, executionController);
Tracing.trace("Acquiring sstable references");
ColumnFamilyStore.ViewFragment view = cfs.select(View.select(SSTableSet.LIVE, partitionKey()));
return queryMemtableAndDiskInternal(cfs, view, null, executionController);
}
private UnfilteredRowIterator queryMemtableAndDiskInternal(ColumnFamilyStore cfs, ReadExecutionController controller)
public UnfilteredRowIterator queryMemtableAndDisk(ColumnFamilyStore cfs,
ColumnFamilyStore.ViewFragment view,
Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> rowTransformer,
ReadExecutionController executionController)
{
assert executionController != null && executionController.validForReadOn(cfs);
Tracing.trace("Executing single-partition query on {}", cfs.name);
return queryMemtableAndDiskInternal(cfs, view, rowTransformer, executionController);
}
private UnfilteredRowIterator queryMemtableAndDiskInternal(ColumnFamilyStore cfs,
ColumnFamilyStore.ViewFragment view,
Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> rowTransformer,
ReadExecutionController controller)
{
/*
* We have 2 main strategies:
@ -685,11 +703,9 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
&& !queriesMulticellType()
&& !controller.isTrackingRepairedStatus())
{
return queryMemtableAndSSTablesInTimestampOrder(cfs, (ClusteringIndexNamesFilter)clusteringIndexFilter(), controller);
return queryMemtableAndSSTablesInTimestampOrder(cfs, view, rowTransformer, (ClusteringIndexNamesFilter)clusteringIndexFilter(), controller);
}
Tracing.trace("Acquiring sstable references");
ColumnFamilyStore.ViewFragment view = cfs.select(View.select(SSTableSet.LIVE, partitionKey()));
view.sstables.sort(SSTableReader.maxTimestampDescending);
ClusteringIndexFilter filter = clusteringIndexFilter();
long minTimestamp = Long.MAX_VALUE;
@ -708,6 +724,9 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
if (memtable.getMinTimestamp() != Memtable.NO_MIN_TIMESTAMP)
minTimestamp = Math.min(minTimestamp, memtable.getMinTimestamp());
if (rowTransformer != null)
iter = Transformation.apply(iter, rowTransformer.apply(memtable));
// Memtable data is always considered unrepaired
controller.updateMinOldestUnrepairedTombstone(memtable.getMinLocalDeletionTime());
inputCollector.addMemtableIterator(RTBoundValidator.validate(iter, RTBoundValidator.Stage.MEMTABLE, false));
@ -767,6 +786,9 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
UnfilteredRowIterator iter = intersects ? makeRowIteratorWithLowerBound(cfs, sstable, metricsCollector)
: makeRowIteratorWithSkippedNonStaticContent(cfs, sstable, metricsCollector);
if (rowTransformer != null)
iter = Transformation.apply(iter, rowTransformer.apply(sstable.getId()));
inputCollector.addSSTableIterator(sstable, iter);
mostRecentPartitionTombstone = Math.max(mostRecentPartitionTombstone,
iter.partitionLevelDeletion().markedForDeleteAt());
@ -789,6 +811,10 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
{
if (!sstable.isRepaired())
controller.updateMinOldestUnrepairedTombstone(sstable.getMinLocalDeletionTime());
if (rowTransformer != null)
iter = Transformation.apply(iter, rowTransformer.apply(sstable.getId()));
inputCollector.addSSTableIterator(sstable, iter);
includedDueToTombstones++;
mostRecentPartitionTombstone = Math.max(mostRecentPartitionTombstone,
@ -922,11 +948,8 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
* no collection or counters are included).
* This method assumes the filter is a {@code ClusteringIndexNamesFilter}.
*/
private UnfilteredRowIterator queryMemtableAndSSTablesInTimestampOrder(ColumnFamilyStore cfs, ClusteringIndexNamesFilter filter, ReadExecutionController controller)
private UnfilteredRowIterator queryMemtableAndSSTablesInTimestampOrder(ColumnFamilyStore cfs, ColumnFamilyStore.ViewFragment view, Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> rowTransformer, ClusteringIndexNamesFilter filter, ReadExecutionController controller)
{
Tracing.trace("Acquiring sstable references");
ColumnFamilyStore.ViewFragment view = cfs.select(View.select(SSTableSet.LIVE, partitionKey()));
ImmutableBTreePartition result = null;
SSTableReadMetricsCollector metricsCollector = new SSTableReadMetricsCollector();
@ -938,7 +961,9 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
if (iter == null)
continue;
result = add(RTBoundValidator.validate(iter, RTBoundValidator.Stage.MEMTABLE, false),
UnfilteredRowIterator wrapped = rowTransformer != null ? Transformation.apply(iter, rowTransformer.apply(memtable))
: iter;
result = add(RTBoundValidator.validate(wrapped, RTBoundValidator.Stage.MEMTABLE, false),
result,
filter,
false,
@ -993,7 +1018,10 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
}
else
{
result = add(RTBoundValidator.validate(iter, RTBoundValidator.Stage.SSTABLE, false),
UnfilteredRowIterator wrapped = rowTransformer != null ? Transformation.apply(iter, rowTransformer.apply(sstable.getId()))
: iter;
result = add(RTBoundValidator.validate(wrapped, RTBoundValidator.Stage.SSTABLE, false),
result,
filter,
sstable.isRepaired(),
@ -1008,8 +1036,9 @@ public class SinglePartitionReadCommand extends ReadCommand implements SinglePar
{
if (iter.isEmpty())
continue;
result = add(RTBoundValidator.validate(iter, RTBoundValidator.Stage.SSTABLE, false),
UnfilteredRowIterator wrapped = rowTransformer != null ? Transformation.apply(iter, rowTransformer.apply(sstable.getId()))
: iter;
result = add(RTBoundValidator.validate(wrapped, RTBoundValidator.Stage.SSTABLE, false),
result,
filter,
sstable.isRepaired(),

6
src/java/org/apache/cassandra/db/lifecycle/Tracker.java
View File

@ -399,7 +399,7 @@ public class Tracker
if (truncating)
notifyRenewed(newMemtable);
else
notifySwitched(result.left.getCurrentMemtable());
notifySwitched(result.left.getCurrentMemtable(), result.right.getCurrentMemtable());
return result.left.getCurrentMemtable();
}
@ -554,9 +554,9 @@ public class Tracker
notify(new MemtableRenewedNotification(renewed));
}
public void notifySwitched(Memtable previous)
public void notifySwitched(Memtable previous, Memtable next)
{
notify(new MemtableSwitchedNotification(previous));
notify(new MemtableSwitchedNotification(previous, next));
}
public void notifyDiscarded(Memtable discarded)

3
src/java/org/apache/cassandra/db/memtable/Memtable.java
View File

@ -22,6 +22,7 @@ import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.concurrent.NotThreadSafe;
import org.apache.cassandra.db.CellSourceIdentifier;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.RegularAndStaticColumns;
@ -54,7 +55,7 @@ import org.apache.cassandra.utils.concurrent.OpOrder;
*
* See Memtable_API.md for details on implementing and using alternative memtable implementations.
*/
public interface Memtable extends Comparable<Memtable>, UnfilteredSource
public interface Memtable extends Comparable<Memtable>, UnfilteredSource, CellSourceIdentifier
{
public static final long NO_MIN_TIMESTAMP = -1;

8
src/java/org/apache/cassandra/index/sai/QueryContext.java
View File

@ -70,8 +70,6 @@ public class QueryContext
* */
public boolean hasUnrepairedMatches = false;
private VectorQueryContext vectorContext;
public QueryContext(ReadCommand readCommand, long executionQuotaMs)
{
this.readCommand = readCommand;
@ -93,10 +91,8 @@ public class QueryContext
}
}
public VectorQueryContext vectorContext()
public int limit()
{
if (vectorContext == null)
vectorContext = new VectorQueryContext(readCommand);
return vectorContext;
return readCommand.limits().count();
}
}

5
src/java/org/apache/cassandra/index/sai/StorageAttachedIndexGroup.java
View File

@ -61,6 +61,7 @@ import org.apache.cassandra.notifications.INotification;
import org.apache.cassandra.notifications.INotificationConsumer;
import org.apache.cassandra.notifications.MemtableDiscardedNotification;
import org.apache.cassandra.notifications.MemtableRenewedNotification;
import org.apache.cassandra.notifications.MemtableSwitchedNotification;
import org.apache.cassandra.notifications.SSTableAddedNotification;
import org.apache.cassandra.notifications.SSTableListChangedNotification;
import org.apache.cassandra.schema.TableMetadata;
@ -275,6 +276,10 @@ public class StorageAttachedIndexGroup implements Index.Group, INotificationCons
{
indexes.forEach(index -> index.memtableIndexManager().renewMemtable(((MemtableRenewedNotification) notification).renewed));
}
else if (notification instanceof MemtableSwitchedNotification)
{
indexes.forEach(index -> index.memtableIndexManager().maybeInitializeMemtableIndex(((MemtableSwitchedNotification) notification).next));
}
else if (notification instanceof MemtableDiscardedNotification)
{
indexes.forEach(index -> index.memtableIndexManager().discardMemtable(((MemtableDiscardedNotification) notification).memtable));

195
src/java/org/apache/cassandra/index/sai/VectorQueryContext.java
View File

@ -1,195 +0,0 @@
/*
* 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.Collection;
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.DiskAnn;
import org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph;
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 final int limit;
// Holds primary keys that are shadowed by expired TTL or row tombstone or range tombstone.
// They are populated by the StorageAttachedIndexSearcher during filtering. They are used to generate
// a bitset for the graph search to indicate graph nodes to ignore.
private TreeSet<PrimaryKey> shadowedPrimaryKeys;
public VectorQueryContext(ReadCommand readCommand)
{
this.limit = readCommand.limits().count();
}
public int limit()
{
return limit;
}
public void recordShadowedPrimaryKeys(Set<PrimaryKey> keys)
{
if (shadowedPrimaryKeys == null)
shadowedPrimaryKeys = new TreeSet<>();
shadowedPrimaryKeys.addAll(keys);
}
// 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 shouldInclude(PrimaryKey pk)
{
return shadowedPrimaryKeys == null || !shadowedPrimaryKeys.contains(pk);
}
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(OnHeapGraph<PrimaryKey> graph)
{
if (shadowedPrimaryKeys == null)
return null;
return new IgnoredKeysBits(graph, shadowedPrimaryKeys);
}
public Bits bitsetForShadowedPrimaryKeys(SegmentMetadata metadata, PrimaryKeyMap primaryKeyMap, DiskAnn 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(ignoredOrdinals, metadata);
}
private static class IgnoringBits implements Bits
{
private final Set<Integer> ignoredOrdinals;
private final int length;
public IgnoringBits(Set<Integer> ignoredOrdinals, SegmentMetadata metadata)
{
this.ignoredOrdinals = ignoredOrdinals;
this.length = 1 + Math.toIntExact(metadata.maxSSTableRowId - metadata.rowIdOffset);
}
@Override
public boolean get(int index)
{
return !ignoredOrdinals.contains(index);
}
@Override
public int length()
{
return length;
}
}
private static class IgnoredKeysBits implements Bits
{
private final OnHeapGraph<PrimaryKey> graph;
private final NavigableSet<PrimaryKey> ignored;
public IgnoredKeysBits(OnHeapGraph<PrimaryKey> graph, NavigableSet<PrimaryKey> ignored)
{
this.graph = graph;
this.ignored = ignored;
}
@Override
public boolean get(int ordinal)
{
Collection<PrimaryKey> keys = graph.keysFromOrdinal(ordinal);
return keys.stream().anyMatch(k -> !ignored.contains(k));
}
@Override
public int length()
{
return graph.size();
}
}
}

121
src/java/org/apache/cassandra/index/sai/disk/EmptyIndex.java Normal file
View File

@ -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;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.virtual.SimpleDataSet;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.SSTableContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
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.utils.CloseableIterator;
/**
* A placeholder index for when there is no on-disk index.
*
* Currenly only used by vector indexes becasue ANN queries require a complete view of the table's sstables, even if
* the associated sstable does not have any data indexed for the column.
*/
public class EmptyIndex extends SSTableIndex
{
public EmptyIndex(SSTableContext sstableContext, StorageAttachedIndex index)
{
super(sstableContext, index);
}
@Override
public long indexFileCacheSize()
{
return 0;
}
@Override
public long getRowCount()
{
return 0;
}
@Override
public long minSSTableRowId()
{
return -1;
}
@Override
public long maxSSTableRowId()
{
return -1;
}
@Override
public ByteBuffer minTerm()
{
return null;
}
@Override
public ByteBuffer maxTerm()
{
return null;
}
@Override
public AbstractBounds<PartitionPosition> bounds()
{
return null;
}
@Override
public List<KeyRangeIterator> search(Expression expression, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
{
return List.of();
}
@Override
public List<CloseableIterator<PrimaryKeyWithScore>> orderBy(Expression orderer, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
{
return List.of();
}
@Override
public List<CloseableIterator<PrimaryKeyWithScore>> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer) throws IOException
{
return List.of();
}
@Override
public void populateSegmentView(SimpleDataSet dataSet)
{
}
@Override
protected void internalRelease()
{
}
}

7
src/java/org/apache/cassandra/index/sai/disk/PrimaryKeyMap.java
View File

@ -26,6 +26,7 @@ import javax.annotation.concurrent.ThreadSafe;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.sstable.SSTableId;
/**
* A bidirectional map of {@link PrimaryKey} to row ID. Implementations of this interface
@ -55,6 +56,12 @@ public interface PrimaryKeyMap extends Closeable
}
}
/**
* Returns the {@link SSTableId} associated with this {@link PrimaryKeyMap}
* @return an {@link SSTableId}
*/
SSTableId getSSTableId();
/**
* Returns a {@link PrimaryKey} for a row ID
*

9
src/java/org/apache/cassandra/index/sai/disk/SSTableIndex.java
View File

@ -37,13 +37,15 @@ import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.SSTableContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.disk.format.Version;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentOrdering;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
import org.apache.cassandra.index.sai.utils.IndexTermType;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.sstable.SSTableIdFactory;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.utils.CloseableIterator;
/**
* A reference-counted container of a {@link SSTableReader} for each column index that:
@ -53,7 +55,7 @@ import org.apache.cassandra.io.sstable.format.SSTableReader;
* <li>Exposes the index metadata for the column index</li>
* </ul>
*/
public abstract class SSTableIndex implements SegmentOrdering, Comparable<SSTableIndex>
public abstract class SSTableIndex implements Comparable<SSTableIndex>
{
private static final Logger logger = LoggerFactory.getLogger(SSTableIndex.class);
@ -142,6 +144,9 @@ public abstract class SSTableIndex implements SegmentOrdering, Comparable<SSTabl
AbstractBounds<PartitionPosition> keyRange,
QueryContext context) throws IOException;
public abstract List<CloseableIterator<PrimaryKeyWithScore>> orderBy(Expression orderer, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException;
public abstract List<CloseableIterator<PrimaryKeyWithScore>> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer) throws IOException;
/**
* Populates a virtual table using the index metadata owned by the index
*/

5
src/java/org/apache/cassandra/index/sai/disk/format/IndexDescriptor.java
View File

@ -35,6 +35,7 @@ import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.IndexValidation;
import org.apache.cassandra.index.sai.SSTableContext;
import org.apache.cassandra.index.sai.disk.EmptyIndex;
import org.apache.cassandra.index.sai.disk.PerColumnIndexWriter;
import org.apache.cassandra.index.sai.disk.PerSSTableIndexWriter;
import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
@ -126,7 +127,9 @@ public class IndexDescriptor
public SSTableIndex newSSTableIndex(SSTableContext sstableContext, StorageAttachedIndex index)
{
return version.onDiskFormat().newSSTableIndex(sstableContext, index);
return isIndexEmpty(index.termType(), index.identifier())
? new EmptyIndex(sstableContext, index)
: version.onDiskFormat().newSSTableIndex(sstableContext, index);
}
public PerSSTableIndexWriter newPerSSTableIndexWriter() throws IOException

17
src/java/org/apache/cassandra/index/sai/disk/v1/SkinnyPrimaryKeyMap.java
View File

@ -34,6 +34,7 @@ import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookupMeta;
import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookup;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.io.util.FileHandle;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.utils.Throwables;
@ -66,6 +67,7 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
protected final LongArray.Factory rowToPartitionReaderFactory;
protected final KeyLookup partitionKeyReader;
protected final PrimaryKey.Factory primaryKeyFactory;
protected final SSTableId sstableId;
private final FileHandle rowToTokenFile;
private final FileHandle rowToPartitionFile;
@ -89,6 +91,7 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
KeyLookupMeta partitionKeysMeta = new KeyLookupMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCKS)));
this.partitionKeyReader = new KeyLookup(partitionKeyBlocksFile, partitionKeyBlockOffsetsFile, partitionKeysMeta, partitionKeyBlockOffsetsMeta);
this.primaryKeyFactory = indexDescriptor.primaryKeyFactory;
this.sstableId = indexDescriptor.sstableDescriptor.id;
}
catch (Throwable t)
{
@ -105,7 +108,8 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
return new SkinnyPrimaryKeyMap(rowIdToToken,
rowIdToPartitionId,
partitionKeyReader.openCursor(),
primaryKeyFactory);
primaryKeyFactory,
sstableId);
}
@Override
@ -119,16 +123,25 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
protected final LongArray rowIdToPartitionIdArray;
protected final KeyLookup.Cursor partitionKeyCursor;
protected final PrimaryKey.Factory primaryKeyFactory;
protected final SSTableId sstableId;
protected SkinnyPrimaryKeyMap(LongArray rowIdToTokenArray,
LongArray rowIdToPartitionIdArray,
KeyLookup.Cursor partitionKeyCursor,
PrimaryKey.Factory primaryKeyFactory)
PrimaryKey.Factory primaryKeyFactory,
SSTableId sstableId)
{
this.rowIdToTokenArray = rowIdToTokenArray;
this.rowIdToPartitionIdArray = rowIdToPartitionIdArray;
this.partitionKeyCursor = partitionKeyCursor;
this.primaryKeyFactory = primaryKeyFactory;
this.sstableId = sstableId;
}
@Override
public SSTableId getSSTableId()
{
return sstableId;
}
@Override

24
src/java/org/apache/cassandra/index/sai/disk/v1/V1SSTableIndex.java
View File

@ -36,12 +36,13 @@ import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.disk.SSTableIndex;
import org.apache.cassandra.index.sai.disk.v1.segment.Segment;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.Throwables;
import static org.apache.cassandra.index.sai.virtual.SegmentsSystemView.CELL_COUNT;
@ -174,14 +175,25 @@ public class V1SSTableIndex extends SSTableIndex
return segmentIterators;
}
@Override
public KeyRangeIterator limitToTopKResults(QueryContext context, List<PrimaryKey> primaryKeys, Expression expression) throws IOException
public List<CloseableIterator<PrimaryKeyWithScore>> orderBy(Expression orderer, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
{
KeyRangeUnionIterator.Builder unionIteratorBuilder = KeyRangeUnionIterator.builder(segments.size());
// Return a list to allow the caller to merge the results from multiple sstables into a single iterator.
List<CloseableIterator<PrimaryKeyWithScore>> iterators = new ArrayList<>(segments.size());
for (Segment segment : segments)
unionIteratorBuilder.add(segment.limitToTopKResults(context, primaryKeys, expression));
if (segment.intersects(keyRange))
iterators.add(segment.orderBy(orderer, keyRange, context));
return unionIteratorBuilder.build();
return iterators;
}
public List<CloseableIterator<PrimaryKeyWithScore>> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer) throws IOException
{
// Return a list to allow the caller to merge the results from multiple sstables into a single iterator.
List<CloseableIterator<PrimaryKeyWithScore>> iterators = new ArrayList<>(segments.size());
for (Segment segment : segments)
iterators.add(segment.orderResultsBy(context, results, orderer));
return iterators;
}
@Override

9
src/java/org/apache/cassandra/index/sai/disk/v1/WidePrimaryKeyMap.java
View File

@ -34,6 +34,7 @@ import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookupMeta;
import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookup;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.io.util.FileHandle;
import org.apache.cassandra.io.util.FileUtils;
@ -103,7 +104,8 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
partitionKeyReader.openCursor(),
clusteringKeyReader.openCursor(),
primaryKeyFactory,
clusteringComparator);
clusteringComparator,
sstableId);
}
@Override
@ -124,9 +126,10 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
KeyLookup.Cursor partitionKeyCursor,
KeyLookup.Cursor clusteringKeyCursor,
PrimaryKey.Factory primaryKeyFactory,
ClusteringComparator clusteringComparator)
ClusteringComparator clusteringComparator,
SSTableId sstableId)
{
super(rowIdToTokenArray, rowIdToPartitionIdArray, partitionKeyCursor, primaryKeyFactory);
super(rowIdToTokenArray, rowIdToPartitionIdArray, partitionKeyCursor, primaryKeyFactory, sstableId);
this.partitionIdToSizeArray = partitionIdToSizeArray;
this.clusteringComparator = clusteringComparator;

75
src/java/org/apache/cassandra/index/sai/disk/v1/postings/VectorPostingList.java
View File

@ -1,75 +0,0 @@
/*
* 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 {@link PostingList} for ANN search results. Transforms result from similarity order to row ID order.
*/
public class VectorPostingList implements PostingList
{
private final LongHeap segmentRowIds;
private final int size;
private final int visitedCount;
public VectorPostingList(PrimitiveIterator.OfInt source, int limit, int visitedCount)
{
this.visitedCount = visitedCount;
segmentRowIds = new LongHeap(Math.max(limit, 1));
int n = 0;
while (source.hasNext() && n++ < limit)
segmentRowIds.push(source.nextInt());
this.size = n;
}
@Override
public long nextPosting()
{
if (segmentRowIds.size() == 0)
return PostingList.END_OF_STREAM;
return segmentRowIds.pop();
}
@Override
public long size()
{
return size;
}
@Override
public long advance(long targetRowID) throws IOException
{
long rowId;
do
{
rowId = nextPosting();
} while (rowId < targetRowID);
return rowId;
}
public int getVisitedCount()
{
return visitedCount;
}
}

21
src/java/org/apache/cassandra/index/sai/disk/v1/segment/IndexSegmentSearcher.java
View File

@ -27,10 +27,13 @@ import org.apache.cassandra.index.sai.StorageAttachedIndex;
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.PostingListRangeIterator;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.postings.PeekablePostingList;
import org.apache.cassandra.index.sai.postings.PostingList;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.utils.CloseableIterator;
/**
* Abstract reader for individual segments of an on-disk index.
@ -57,12 +60,13 @@ public abstract class IndexSegmentSearcher implements SegmentOrdering, Closeable
}
public static IndexSegmentSearcher open(PrimaryKeyMap.Factory primaryKeyMapFactory,
SSTableId sstableId,
PerColumnIndexFiles indexFiles,
SegmentMetadata segmentMetadata,
StorageAttachedIndex index) throws IOException
{
if (index.termType().isVector())
return new VectorIndexSegmentSearcher(primaryKeyMapFactory, indexFiles, segmentMetadata, index);
return new VectorIndexSegmentSearcher(primaryKeyMapFactory, sstableId, indexFiles, segmentMetadata, index);
else if (index.termType().isLiteral())
return new LiteralIndexSegmentSearcher(primaryKeyMapFactory, indexFiles, segmentMetadata, index);
else
@ -84,6 +88,21 @@ public abstract class IndexSegmentSearcher implements SegmentOrdering, Closeable
*/
public abstract KeyRangeIterator search(Expression expression, AbstractBounds<PartitionPosition> keyRange, QueryContext queryContext) throws IOException;
/**
* Order the rows by the given expression.
*
* @param orderer the object containing the ordering logic
* @param keyRange key range specific in read command, used by ANN index
* @param context to track per sstable cache and per query metrics
*
* @return an iterator of {@link PrimaryKeyWithScore} in descending score order
*/
public CloseableIterator<PrimaryKeyWithScore> orderBy(Expression orderer, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
{
throw new UnsupportedOperationException();
}
KeyRangeIterator toPrimaryKeyIterator(PostingList postingList, QueryContext queryContext) throws IOException
{
if (postingList == null || postingList.size() == 0)

28
src/java/org/apache/cassandra/index/sai/disk/v1/segment/Segment.java
View File

@ -30,12 +30,13 @@ import org.apache.cassandra.dht.Token;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.SSTableContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
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.PrimaryKeyWithScore;
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.io.util.FileUtils;
import org.apache.cassandra.utils.CloseableIterator;
/**
* Each segment represents an on-disk index structure (balanced tree/terms/postings) flushed by memory limit or token boundaries.
@ -47,8 +48,6 @@ public class Segment implements SegmentOrdering, Closeable
private final Token.KeyBound minKeyBound;
private final Token.KeyBound maxKeyBound;
// per sstable
final PrimaryKeyMap.Factory primaryKeyMapFactory;
// per-segment
public final SegmentMetadata metadata;
@ -59,16 +58,14 @@ public class Segment implements SegmentOrdering, Closeable
this.minKeyBound = metadata.minKey.token().minKeyBound();
this.maxKeyBound = metadata.maxKey.token().maxKeyBound();
this.primaryKeyMapFactory = sstableContext.primaryKeyMapFactory;
this.metadata = metadata;
this.index = IndexSegmentSearcher.open(primaryKeyMapFactory, indexFiles, metadata, index);
this.index = IndexSegmentSearcher.open(sstableContext.primaryKeyMapFactory, sstableContext.sstable.getId(), indexFiles, metadata, index);
}
@VisibleForTesting
public Segment(Token minKey, Token maxKey)
{
this.primaryKeyMapFactory = null;
this.metadata = null;
this.minKeyBound = minKey.minKeyBound();
this.maxKeyBound = maxKey.maxKeyBound();
@ -112,10 +109,23 @@ public class Segment implements SegmentOrdering, Closeable
return index.search(expression, keyRange, context);
}
@Override
public KeyRangeIterator limitToTopKResults(QueryContext context, List<PrimaryKey> primaryKeys, Expression expression) throws IOException
/**
* Order the on-disk index synchronously and produce an iterator in score order
*
* @param orderer the expression to use when searching the on disk index
* @param keyRange key range specific in read command, used by ANN index
* @param context to track per sstable cache and per query metrics
* @return an iterator of {@link PrimaryKeyWithScore} in score order
*/
public CloseableIterator<PrimaryKeyWithScore> orderBy(Expression orderer, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
{
return index.limitToTopKResults(context, primaryKeys, expression);
return index.orderBy(orderer, keyRange, context);
}
@Override
public CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer) throws IOException
{
return index.orderResultsBy(context, results, orderer);
}
@Override

2
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentBuilder.java
View File

@ -113,7 +113,7 @@ public abstract class SegmentBuilder
public VectorSegmentBuilder(StorageAttachedIndex index, NamedMemoryLimiter limiter)
{
super(index, limiter);
graphIndex = new OnHeapGraph<>(index.termType().indexType(), index.indexWriterConfig(), false);
graphIndex = new OnHeapGraph<>(index.termType().indexType(), index.indexWriterConfig(), null);
}
@Override

22
src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java
View File

@ -23,24 +23,28 @@ import java.util.List;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
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.utils.CloseableIterator;
/**
* A {@link SegmentOrdering} orders and limits a list of {@link PrimaryKey}s.
* <p>
* A {@link SegmentOrdering} orders an index and produces a stream of {@link PrimaryKeyWithScore}s.
*
* The limit can be used to lazily order the {@link PrimaryKey}s. Due to the possiblity for
* shadowed or updated keys, a {@link SegmentOrdering} should be able to order the whole index
* until exhausted.
*
* When using {@link SegmentOrdering} there are several steps to
* build the list of Primary Keys to be ordered and limited:
* <p>
* build the list of Primary Keys to be ordered:
*
* 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.
* 3. Fan the primary keys from step 2 out to each sstable segment to order the list of primary keys.
* <p>
* SegmentOrdering handles the fourth step.
* SegmentOrdering handles the third 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
@ -51,7 +55,7 @@ public interface SegmentOrdering
/**
* Reorder, limit, and put back into original order the results from a single sstable
*/
default KeyRangeIterator limitToTopKResults(QueryContext queryContext, List<PrimaryKey> primaryKeys, Expression expression) throws IOException
default CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext queryContext, List<PrimaryKey> results, Expression orderer) throws IOException
{
throw new UnsupportedOperationException();
}

307
src/java/org/apache/cassandra/index/sai/disk/v1/segment/VectorIndexSegmentSearcher.java
View File

@ -20,37 +20,47 @@ package org.apache.cassandra.index.sai.disk.v1.segment;
import java.io.IOException;
import java.lang.invoke.MethodHandles;
import java.util.List;
import java.util.function.IntConsumer;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import com.google.common.annotations.VisibleForTesting;
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.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
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.VectorPostingList;
import org.apache.cassandra.index.sai.disk.v1.vector.BruteForceRowIdIterator;
import org.apache.cassandra.index.sai.disk.v1.vector.DiskAnn;
import org.apache.cassandra.index.sai.disk.v1.vector.NeighborQueueRowIdIterator;
import org.apache.cassandra.index.sai.disk.v1.vector.OnDiskOrdinalsMap;
import org.apache.cassandra.index.sai.disk.v1.vector.OptimizeFor;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.disk.v1.vector.RowIdToPrimaryKeyWithScoreIterator;
import org.apache.cassandra.index.sai.disk.v1.vector.RowIdWithScore;
import org.apache.cassandra.index.sai.disk.v1.vector.SegmentRowIdOrdinalPairs;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.iterators.KeyRangeListIterator;
import org.apache.cassandra.index.sai.memory.VectorMemoryIndex;
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.AtomicRatio;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.RangeUtil;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.tracing.Tracing;
import org.apache.cassandra.utils.CloseableIterator;
import io.github.jbellis.jvector.graph.GraphIndex;
import io.github.jbellis.jvector.graph.NeighborQueue;
import io.github.jbellis.jvector.graph.NeighborSimilarity;
import io.github.jbellis.jvector.pq.CompressedVectors;
import io.github.jbellis.jvector.util.Bits;
import io.github.jbellis.jvector.util.SparseFixedBitSet;
import static java.lang.Math.max;
import static java.lang.Math.min;
@ -62,22 +72,27 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
{
private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
// If true, use brute force. If false, use graph search. If null, use the normal logic.
@VisibleForTesting
public static Boolean FORCE_BRUTE_FORCE_ANN = null;
private final DiskAnn graph;
private final int globalBruteForceRows;
private final AtomicRatio actualExpectedRatio = new AtomicRatio();
private final ThreadLocal<SparseFixedBitSet> cachedBitSets;
private final OptimizeFor optimizeFor;
private final ColumnMetadata column;
VectorIndexSegmentSearcher(PrimaryKeyMap.Factory primaryKeyMapFactory,
SSTableId sstableId,
PerColumnIndexFiles perIndexFiles,
SegmentMetadata segmentMetadata,
StorageAttachedIndex index) throws IOException
{
super(primaryKeyMapFactory, perIndexFiles, segmentMetadata, index);
graph = new DiskAnn(segmentMetadata.componentMetadatas, perIndexFiles, index.indexWriterConfig());
graph = new DiskAnn(segmentMetadata.componentMetadatas, perIndexFiles, index.indexWriterConfig(), sstableId);
cachedBitSets = ThreadLocal.withInitial(() -> new SparseFixedBitSet(graph.size()));
globalBruteForceRows = Integer.MAX_VALUE;
optimizeFor = index.indexWriterConfig().getOptimizeFor();
column = index.termType().columnMetadata();
}
@Override
@ -87,74 +102,77 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
}
@Override
public KeyRangeIterator search(Expression exp, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
public KeyRangeIterator search(Expression expression, AbstractBounds<PartitionPosition> keyRange, QueryContext queryContext) throws IOException
{
int limit = context.vectorContext().limit();
if (logger.isTraceEnabled())
logger.trace(index.identifier().logMessage("Searching on expression '{}'..."), exp);
if (exp.getIndexOperator() != Expression.IndexOperator.ANN)
throw new IllegalArgumentException(index.identifier().logMessage("Unsupported expression during ANN index query: " + exp));
int topK = optimizeFor.topKFor(limit);
BitsOrPostingList bitsOrPostingList = bitsOrPostingListForKeyRange(context.vectorContext(), keyRange, topK);
if (bitsOrPostingList.skipANN())
return toPrimaryKeyIterator(bitsOrPostingList.postingList(), context);
float[] queryVector = index.termType().decomposeVector(exp.lower().value.raw.duplicate());
VectorPostingList vectorPostings = graph.search(queryVector, topK, limit, bitsOrPostingList.getBits());
if (bitsOrPostingList.expectedNodesVisited >= 0)
updateExpectedNodes(vectorPostings.getVisitedCount(), bitsOrPostingList.expectedNodesVisited);
return toPrimaryKeyIterator(vectorPostings, context);
throw new UnsupportedOperationException();
}
/**
* 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, int limit) throws IOException
@Override
public CloseableIterator<PrimaryKeyWithScore> orderBy(Expression orderer, AbstractBounds<PartitionPosition> keyRange, QueryContext context) throws IOException
{
int limit = context.limit();
if (logger.isTraceEnabled())
logger.trace(index.identifier().logMessage("Searching on expression '{}'..."), orderer);
if (orderer.getIndexOperator() != Expression.IndexOperator.ANN)
throw new IllegalArgumentException(index.identifier().logMessage("Unsupported expression during ANN index query: " + orderer));
int topK = optimizeFor.topKFor(limit);
float[] queryVector = index.termType().decomposeVector(orderer.lower().value.raw.duplicate());
CloseableIterator<RowIdWithScore> result = searchInternal(keyRange, queryVector, limit, topK);
return toScoreSortedIterator(result);
}
private CloseableIterator<RowIdWithScore> searchInternal(AbstractBounds<PartitionPosition> keyRange, float[] queryVector, int limit, int topK) throws IOException
{
try (PrimaryKeyMap primaryKeyMap = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap())
{
// not restricted
if (RangeUtil.coversFullRing(keyRange))
return new BitsOrPostingList(context.bitsetForShadowedPrimaryKeys(metadata, primaryKeyMap, graph));
return searchInternalUnrestricted(queryVector, limit, topK);
// it will return the next row id if given key is not found.
long minSSTableRowId = primaryKeyMap.ceiling(keyRange.left.getToken());
// If we didn't find the first key, we won't find the last primary key either
if (minSSTableRowId < 0)
return new BitsOrPostingList(PostingList.EMPTY);
return CloseableIterator.empty();
long maxSSTableRowId = getMaxSSTableRowId(primaryKeyMap, keyRange.right);
if (minSSTableRowId > maxSSTableRowId)
return new BitsOrPostingList(PostingList.EMPTY);
return CloseableIterator.empty();
// if it covers entire segment, skip bit set
if (minSSTableRowId <= metadata.minSSTableRowId && maxSSTableRowId >= metadata.maxSSTableRowId)
return new BitsOrPostingList(context.bitsetForShadowedPrimaryKeys(metadata, primaryKeyMap, graph));
return searchInternalUnrestricted(queryVector, limit, topK);
minSSTableRowId = Math.max(minSSTableRowId, metadata.minSSTableRowId);
maxSSTableRowId = min(maxSSTableRowId, metadata.maxSSTableRowId);
// If num of matches are not bigger than limit, skip ANN.
// (nRows should not include shadowed rows, but context doesn't break those out by segment,
// so we will live with the inaccuracy.)
// If num of matches are not bigger than limit, skip graph search and lazily sort by brute force.
int nRows = Math.toIntExact(maxSSTableRowId - minSSTableRowId + 1);
int maxBruteForceRows = min(globalBruteForceRows, maxBruteForceRows(limit, nRows, graph.size()));
int maxBruteForceRows = maxBruteForceRows(limit, nRows, graph.size());
logger.trace("Search range covers {} rows; max brute force rows is {} for sstable index with {} nodes, LIMIT {}",
nRows, maxBruteForceRows, graph.size(), limit);
Tracing.trace("Search range covers {} rows; max brute force rows is {} for sstable index with {} nodes, LIMIT {}",
nRows, maxBruteForceRows, graph.size(), limit);
if (nRows <= maxBruteForceRows)
boolean shouldBruteForce = FORCE_BRUTE_FORCE_ANN == null ? nRows <= maxBruteForceRows : FORCE_BRUTE_FORCE_ANN;
if (shouldBruteForce)
{
IntArrayList postings = new IntArrayList(Math.toIntExact(nRows), -1);
for (long sstableRowId = minSSTableRowId; sstableRowId <= maxSSTableRowId; sstableRowId++)
SegmentRowIdOrdinalPairs segmentOrdinalPairs = new SegmentRowIdOrdinalPairs(Math.toIntExact(nRows));
try (OnDiskOrdinalsMap.OrdinalsView ordinalsView = graph.getOrdinalsView())
{
if (context.shouldInclude(sstableRowId, primaryKeyMap))
postings.addInt(metadata.toSegmentRowId(sstableRowId));
for (long sstableRowId = minSSTableRowId; sstableRowId <= maxSSTableRowId; sstableRowId++)
{
int segmentRowId = metadata.toSegmentRowId(sstableRowId);
int ordinal = ordinalsView.getOrdinalForRowId(segmentRowId);
if (ordinal >= 0)
segmentOrdinalPairs.add(segmentRowId, ordinal);
}
}
return new BitsOrPostingList(new IntArrayPostingList(postings.toIntArray()));
return orderByBruteForce(queryVector, segmentOrdinalPairs, limit, topK);
}
// create a bitset of ordinals corresponding to the rows in the given key range
@ -164,15 +182,12 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
{
for (long sstableRowId = minSSTableRowId; sstableRowId <= maxSSTableRowId; sstableRowId++)
{
if (context.shouldInclude(sstableRowId, primaryKeyMap))
int segmentRowId = metadata.toSegmentRowId(sstableRowId);
int ordinal = ordinalsView.getOrdinalForRowId(segmentRowId);
if (ordinal >= 0)
{
int segmentRowId = metadata.toSegmentRowId(sstableRowId);
int ordinal = ordinalsView.getOrdinalForRowId(segmentRowId);
if (ordinal >= 0)
{
bits.set(ordinal);
hasMatches = true;
}
bits.set(ordinal);
hasMatches = true;
}
}
}
@ -182,12 +197,21 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
}
if (!hasMatches)
return new BitsOrPostingList(PostingList.EMPTY);
return CloseableIterator.empty();
return new BitsOrPostingList(bits, VectorMemoryIndex.expectedNodesVisited(limit, nRows, graph.size()));
int expectedNodesVisited = expectedNodesVisited(limit, bits.cardinality(), graph.size());
IntConsumer nodesVisitedConsumer = nodesVisited -> updateExpectedNodes(nodesVisited, expectedNodesVisited);
return graph.search(queryVector, topK, limit, bits, nodesVisitedConsumer);
}
}
private CloseableIterator<RowIdWithScore> searchInternalUnrestricted(float[] queryVector, int limit, int topK)
{
int expectedNodesVisited = expectedNodesVisited(limit, graph.size(), graph.size());
IntConsumer nodesVisitedConsumer = nodesVisited -> updateExpectedNodes(nodesVisited, expectedNodesVisited);
return graph.search(queryVector, topK, limit, new Bits.MatchAllBits(graph.size()), nodesVisitedConsumer);
}
private long getMaxSSTableRowId(PrimaryKeyMap primaryKeyMap, PartitionPosition right)
{
// if the right token is the minimum token, there is no upper bound on the keyRange and
@ -208,29 +232,72 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
return bits;
}
@Override
public KeyRangeIterator limitToTopKResults(QueryContext context, List<PrimaryKey> primaryKeys, Expression expression) throws IOException
/**
* Produces a descending score ordered iterator over the rows in the given segment. Branches depending on the number
* of rows to consider and whether the graph has compressed vectors available for faster comparisons.
*/
private CloseableIterator<RowIdWithScore> orderByBruteForce(float[] queryVector, SegmentRowIdOrdinalPairs segmentOrdinalPairs, int limit, int topK) throws IOException
{
int limit = context.vectorContext().limit();
if (segmentOrdinalPairs.size() == 0)
return CloseableIterator.empty();
// If we have more than topK segmentOrdinalPairs, we do a two pass partial sort by first getting the approximate
// similarity score via the PQ vectors that are already in memory and then by hitting disk to get the full
// precision vectors to get the full precision similarity score.
if (graph.getCompressedVectors() != null && segmentOrdinalPairs.size() > topK)
return orderByBruteForceTwoPass(graph.getCompressedVectors(), queryVector, segmentOrdinalPairs, limit, topK);
try (GraphIndex.View<float[]> view = graph.getView())
{
NeighborSimilarity.ExactScoreFunction esf = graph.getExactScoreFunction(queryVector, view);
NeighborQueue scoredRowIds = segmentOrdinalPairs.mapToSegmentRowIdScoreHeap(esf);
return new NeighborQueueRowIdIterator(scoredRowIds);
}
catch (Exception e)
{
throw new IOException(e);
}
}
/**
* Materialize the compressed vectors for the given segment row ids, put them into a priority queue ordered by
* approximate similarity score, and then pass to the {@link BruteForceRowIdIterator} to lazily resolve the
* full resolution ordering as needed.
*/
private CloseableIterator<RowIdWithScore> orderByBruteForceTwoPass(CompressedVectors cv,
float[] queryVector,
SegmentRowIdOrdinalPairs segmentOrdinalPairs,
int limit,
int rerankK)
{
NeighborSimilarity.ApproximateScoreFunction scoreFunction = graph.getApproximateScoreFunction(queryVector);
// Store the index of the (rowId, ordinal) pair from the segmentOrdinalPairs in the NodeQueue so that we can
// retrieve both values with O(1) lookup when we need to resolve the full resolution score in the
// BruteForceRowIdIterator.
NeighborQueue approximateScoreHeap = segmentOrdinalPairs.mapToIndexScoreIterator(scoreFunction);
GraphIndex.View<float[]> view = graph.getView();
NeighborSimilarity.ExactScoreFunction esf = graph.getExactScoreFunction(queryVector, view);
return new BruteForceRowIdIterator(approximateScoreHeap, segmentOrdinalPairs, esf, limit, rerankK, view);
}
@Override
public CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer) throws IOException
{
int limit = context.limit();
// VSTODO would it be better to do a binary search to find the boundaries?
List<PrimaryKey> keysInRange = primaryKeys.stream()
.dropWhile(k -> k.compareTo(metadata.minKey) < 0)
.takeWhile(k -> k.compareTo(metadata.maxKey) <= 0)
.collect(Collectors.toList());
List<PrimaryKey> keysInRange = results.stream()
.dropWhile(k -> k.compareTo(metadata.minKey) < 0)
.takeWhile(k -> k.compareTo(metadata.maxKey) <= 0)
.collect(Collectors.toList());
if (keysInRange.isEmpty())
return KeyRangeIterator.empty();
int topK = optimizeFor.topKFor(limit);
if (shouldUseBruteForce(topK, limit, keysInRange.size()))
return new KeyRangeListIterator(metadata.minKey, metadata.maxKey, keysInRange);
return CloseableIterator.empty();
try (PrimaryKeyMap primaryKeyMap = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap())
{
// the iterator represents keys from the whole table -- we'll only pull of those that
// are from our own token range, so we can use row ids to order the results by vector similarity.
int maxSegmentRowId = metadata.toSegmentRowId(metadata.maxSSTableRowId);
SparseFixedBitSet bits = bitSetForSearch();
IntArrayList rowIds = new IntArrayList();
try (var ordinalsView = graph.getOrdinalsView())
SegmentRowIdOrdinalPairs segmentOrdinalPairs = new SegmentRowIdOrdinalPairs(keysInRange.size());
try (OnDiskOrdinalsMap.OrdinalsView ordinalsView = graph.getOrdinalsView())
{
for (PrimaryKey primaryKey : keysInRange)
{
@ -245,45 +312,56 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
break;
int segmentRowId = metadata.toSegmentRowId(sstableRowId);
rowIds.add(segmentRowId);
// VSTODO now that we know the size of keys evaluated, is it worth doing the brute
// force check eagerly to potentially skip the PK to sstable row id to ordinal lookup?
int ordinal = ordinalsView.getOrdinalForRowId(segmentRowId);
if (ordinal >= 0)
bits.set(ordinal);
segmentOrdinalPairs.add(segmentRowId, ordinal);
}
}
if (shouldUseBruteForce(topK, limit, rowIds.size()))
return toPrimaryKeyIterator(new IntArrayPostingList(rowIds.toIntArray()), context);
int topK = optimizeFor.topKFor(limit);
float[] queryVector = index.termType().decomposeVector(orderer.lower().value.raw.duplicate());
if (shouldUseBruteForce(topK, limit, segmentOrdinalPairs.size()))
{
return toScoreSortedIterator(orderByBruteForce(queryVector, segmentOrdinalPairs, limit, topK));
}
SparseFixedBitSet bits = bitSetForSearch();
segmentOrdinalPairs.forEachOrdinal(bits::set);
// else ask the index to perform a search limited to the bits we created
float[] queryVector = index.termType().decomposeVector(expression.lower().value.raw.duplicate());
VectorPostingList results = graph.search(queryVector, topK, limit, bits);
updateExpectedNodes(results.getVisitedCount(), expectedNodesVisited(topK, maxSegmentRowId, graph.size()));
return toPrimaryKeyIterator(results, context);
int expectedNodesVisited = expectedNodesVisited(limit, segmentOrdinalPairs.size(), graph.size());
IntConsumer nodesVisitedConsumer = nodesVisited -> updateExpectedNodes(nodesVisited, expectedNodesVisited);
CloseableIterator<RowIdWithScore> result = graph.search(queryVector, topK, limit, bits, nodesVisitedConsumer);
return toScoreSortedIterator(result);
}
}
private boolean shouldUseBruteForce(int topK, int limit, int numRows)
{
// if we have a small number of results then let TopK processor do exact NN computation
int maxBruteForceRows = min(globalBruteForceRows, maxBruteForceRows(topK, numRows, graph.size()));
int maxBruteForceRows = maxBruteForceRows(topK, numRows, graph.size());
logger.trace("SAI materialized {} rows; max brute force rows is {} for sstable index with {} nodes, LIMIT {}",
numRows, maxBruteForceRows, graph.size(), limit);
Tracing.trace("SAI materialized {} rows; max brute force rows is {} for sstable index with {} nodes, LIMIT {}",
numRows, maxBruteForceRows, graph.size(), limit);
return numRows <= maxBruteForceRows;
return FORCE_BRUTE_FORCE_ANN == null ? numRows <= maxBruteForceRows
: FORCE_BRUTE_FORCE_ANN;
}
private int maxBruteForceRows(int limit, int nPermittedOrdinals, int graphSize)
{
int expectedNodes = expectedNodesVisited(limit, nPermittedOrdinals, graphSize);
// ANN index will do a bunch of extra work besides the full comparisons (performing PQ similarity for each edge);
// brute force from sstable will also do a bunch of extra work (going through trie index to look up row).
// VSTODO I'm not sure which one is more expensive (and it depends on things like sstable chunk cache hit ratio)
// so I'm leaving it as a 1:1 ratio for now.
return max(limit, expectedNodes);
int expectedNodesVisited = expectedNodesVisited(limit, nPermittedOrdinals, graphSize);
int expectedComparisons = index.indexWriterConfig().getMaximumNodeConnections() * expectedNodesVisited;
// Brute force here means reading each vector from the index file on disk, comparing the row's vector to the
// search vector to get a score, and then putting the results into a priority queue to then iterate over.
// Alternatively, we search the graph, which entails comparisons and disk reads. The goal is to reduce disk
// accesses and number of comparisons.
// VSTODO the below factor is dramatically oversimplified
// larger dimension should increase this, because comparisons are more expensive
double memoryToDiskFactor = 0.25;
return (int) max(limit, memoryToDiskFactor * expectedComparisons);
}
private int expectedNodesVisited(int limit, int nPermittedOrdinals, int graphSize)
@ -297,7 +375,7 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
assert expectedNodesVisited >= 0 : expectedNodesVisited;
assert actualNodesVisited >= 0 : actualNodesVisited;
if (actualNodesVisited >= 1000 && actualNodesVisited > 2 * expectedNodesVisited || expectedNodesVisited > 2 * actualNodesVisited)
logger.warn("Predicted visiting {} nodes, but actually visited {}", expectedNodesVisited, actualNodesVisited);
logger.trace("Predicted visiting {} nodes, but actually visited {}", expectedNodesVisited, actualNodesVisited);
actualExpectedRatio.update(actualNodesVisited, expectedNodesVisited);
}
@ -313,49 +391,14 @@ public class VectorIndexSegmentSearcher extends IndexSegmentSearcher
graph.close();
}
private static class BitsOrPostingList
private CloseableIterator<PrimaryKeyWithScore> toScoreSortedIterator(CloseableIterator<RowIdWithScore> rowIdIterator) throws IOException
{
private final Bits bits;
private final int expectedNodesVisited;
private final PostingList postingList;
public BitsOrPostingList(@Nullable Bits bits, int expectedNodesVisited)
if (!rowIdIterator.hasNext())
{
this.bits = bits;
this.expectedNodesVisited = expectedNodesVisited;
this.postingList = null;
FileUtils.closeQuietly(rowIdIterator);
return CloseableIterator.empty();
}
public BitsOrPostingList(@Nullable Bits bits)
{
this.bits = bits;
this.postingList = null;
this.expectedNodesVisited = -1;
}
public BitsOrPostingList(PostingList postingList)
{
this.bits = null;
this.postingList = Preconditions.checkNotNull(postingList);
this.expectedNodesVisited = -1;
}
@Nullable
public Bits getBits()
{
Preconditions.checkState(!skipANN());
return bits;
}
public PostingList postingList()
{
Preconditions.checkState(skipANN());
return postingList;
}
public boolean skipANN()
{
return postingList != null;
}
return new RowIdToPrimaryKeyWithScoreIterator(column, primaryKeyMapFactory, rowIdIterator, metadata.rowIdOffset);
}
}

162
src/java/org/apache/cassandra/index/sai/disk/v1/vector/AutoResumingNodeScoreIterator.java Normal file
View File

@ -0,0 +1,162 @@
/*
* 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.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.function.IntConsumer;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.tracing.Tracing;
import org.apache.cassandra.utils.AbstractIterator;
import io.github.jbellis.jvector.graph.GraphIndex;
import io.github.jbellis.jvector.graph.GraphSearcher;
import io.github.jbellis.jvector.graph.NeighborSimilarity;
import io.github.jbellis.jvector.graph.SearchResult;
import io.github.jbellis.jvector.util.Bits;
import io.github.jbellis.jvector.util.GrowableBitSet;
/**
* An iterator over {@link SearchResult.NodeScore} backed by a {@link SearchResult} that resumes search
* when the backing {@link SearchResult} is exhausted.
*/
public class AutoResumingNodeScoreIterator extends AbstractIterator<SearchResult.NodeScore>
{
private final GraphSearcher<float[]> searcher;
private final GraphIndex.View<float[]> view;
private final NeighborSimilarity.ScoreFunction scoreFunction;
private final NeighborSimilarity.ReRanker<float[]> reRanker;
private final int topK;
private final Bits acceptBits;
private final boolean inMemory;
private final String source;
private final IntConsumer nodesVisitedConsumer;
private Iterator<SearchResult.NodeScore> nodeScores = Collections.emptyIterator();
private int cumulativeNodesVisited;
// Defer initialization since it is only needed if we need to resume search
private SkipVisitedBits visited = null;
private SearchResult.NodeScore[] previousResult = null;
/**
* Create a new {@link AutoResumingNodeScoreIterator} that iterates over the provided {@link SearchResult}.
* If the {@link SearchResult} is consumed, it retrieves the next {@link SearchResult} until the search returns
* no more results.
* @param searcher the {@link GraphSearcher} to use to search and resume search.
* @param nodesVisitedConsumer a consumer that accepts the total number of nodes visited
* @param inMemory whether the graph is in memory or on disk (used for trace logging)
* @param source the source of the search (used for trace logging)
* @param view the view used to read from disk. It will be closed when the iterator is closed.
*/
public AutoResumingNodeScoreIterator(GraphSearcher<float[]> searcher,
NeighborSimilarity.ScoreFunction scoreFunction,
NeighborSimilarity.ReRanker<float[]> reRanker,
int topK,
Bits acceptBits,
IntConsumer nodesVisitedConsumer,
boolean inMemory,
String source,
GraphIndex.View<float[]> view)
{
this.searcher = searcher;
this.scoreFunction = scoreFunction;
this.reRanker = reRanker;
this.topK = topK;
this.acceptBits = acceptBits;
this.cumulativeNodesVisited = 0;
this.nodesVisitedConsumer = nodesVisitedConsumer;
this.inMemory = inMemory;
this.source = source;
this.view = view;
}
@Override
protected SearchResult.NodeScore computeNext()
{
if (nodeScores.hasNext())
return nodeScores.next();
// Add result from previous search to visited bits
if (previousResult != null)
{
if (visited == null)
visited = new SkipVisitedBits(acceptBits, previousResult.length);
visited.visited(previousResult);
}
Bits bits = visited == null ? acceptBits : visited;
SearchResult nextResult = searcher.search(scoreFunction, reRanker, topK, bits);
// Record metrics (we add here instead of overwriting because re-queries are expensive proportional to the
// number of visited nodes and even though we throw away some of those results, it helps us determine the
// right path for brute force vs. ANN)
cumulativeNodesVisited += nextResult.getVisitedCount();
if (Tracing.isTracing())
{
Tracing.trace("{} based ANN {} for topK {} visited {} nodes to return {} results from {}",
inMemory ? "Memory" : "Disk", previousResult == null ? "initial" : "re-query",
topK, nextResult.getVisitedCount(), nextResult.getNodes().length, source);
}
previousResult = nextResult.getNodes();
// If the next result is empty, we are done searching.
nodeScores = Arrays.stream(nextResult.getNodes()).iterator();
return nodeScores.hasNext() ? nodeScores.next() : endOfData();
}
@Override
public void close()
{
nodesVisitedConsumer.accept(cumulativeNodesVisited);
FileUtils.closeQuietly(view);
}
private static class SkipVisitedBits implements Bits
{
private final Bits acceptBits;
private final GrowableBitSet visited;
SkipVisitedBits(Bits acceptBits, int initialBits)
{
this.acceptBits = acceptBits;
this.visited = new GrowableBitSet(initialBits);
}
void visited(SearchResult.NodeScore[] nodes)
{
for (SearchResult.NodeScore nodeScore : nodes)
visited.set(nodeScore.node);
}
@Override
public boolean get(int i)
{
return acceptBits.get(i) && !visited.get(i);
}
@Override
public int length()
{
return acceptBits.length();
}
}
}

124
src/java/org/apache/cassandra/index/sai/disk/v1/vector/BruteForceRowIdIterator.java Normal file
View File

@ -0,0 +1,124 @@
/*
* 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 javax.annotation.concurrent.NotThreadSafe;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.utils.AbstractIterator;
import io.github.jbellis.jvector.graph.GraphIndex;
import io.github.jbellis.jvector.graph.NeighborQueue;
import io.github.jbellis.jvector.graph.NeighborSimilarity;
/**
* An iterator over {@link RowIdWithScore} that lazily consumes from a {@link NeighborQueue} of approximate scores.
* <p>
* The idea is that we maintain the same level of accuracy as we would get from a graph search, by re-ranking the top
* `k` best approximate scores at a time with the full resolution vectors to return the top `limit`.
* <p>
* For example, suppose that limit=3 and k=5 and we have ten elements. After our first re-ranking batch, we have
* ABDEF?????
* We will return A, B, and D; if more elements are requested, we will re-rank another 5 (so three more, including
* the two remaining from the first batch). Here we uncover C, G, and H, and order them appropriately:
* CEFGH??
* This illustrates that, also like a graph search, we only guarantee ordering of results within a re-ranking batch,
* not globally.
* <p>
* Note that we deliberately do not fetch new items from the approximate list until the first batch of `limit`-many
* is consumed. We do this because we expect that most often the first limit-many will pass the final verification
* and only query more if some didn't (e.g. because the vector was deleted in a newer sstable).
* <p>
* As an implementation detail, we use a heap to maintain state rather than a List and sorting.
*/
@NotThreadSafe
public class BruteForceRowIdIterator extends AbstractIterator<RowIdWithScore>
{
// We use two binary heaps (NeighborQueue) because we do not need an eager ordering of
// these results. Depending on how many sstables the query hits and the relative scores of vectors from those
// sstables, we may not need to return more than the first handful of scores.
// Heap with compressed vector scores
private final NeighborQueue approximateScoreQueue;
private final SegmentRowIdOrdinalPairs segmentOrdinalPairs;
// Use the jvector NeighborQueue to avoid unnecessary object allocations
private final NeighborQueue exactScoreQueue;
private final NeighborSimilarity.ExactScoreFunction reranker;
private final GraphIndex.View<float[]> view;
private final int topK;
private final int limit;
private int rerankedCount;
/**
* @param approximateScoreQueue A heap of indexes ordered by their approximate similarity scores
* @param segmentOrdinalPairs A mapping from the index in the approximateScoreQueue to the node's rowId and ordinal
* @param reranker A function that takes a graph ordinal and returns the exact similarity score
* @param limit The query limit
* @param topK The number of vectors to resolve and score before returning results
* @param view The view of the graph, passed so we can close it when the iterator is closed
*/
public BruteForceRowIdIterator(NeighborQueue approximateScoreQueue,
SegmentRowIdOrdinalPairs segmentOrdinalPairs,
NeighborSimilarity.ExactScoreFunction reranker,
int limit,
int topK,
GraphIndex.View<float[]> view)
{
this.approximateScoreQueue = approximateScoreQueue;
this.segmentOrdinalPairs = segmentOrdinalPairs;
this.exactScoreQueue = new NeighborQueue(limit, true);
this.reranker = reranker;
assert topK >= limit : "topK must be greater than or equal to limit. Found: " + topK + " < " + limit;
this.limit = limit;
this.topK = topK;
this.rerankedCount = topK; // placeholder to kick off computeNext
this.view = view;
}
@Override
protected RowIdWithScore computeNext()
{
int consumed = rerankedCount - exactScoreQueue.size();
if (consumed >= limit)
{
// Refill the exactScoreQueue until it reaches topK exact scores, or the approximate score queue is empty
while (approximateScoreQueue.size() > 0 && exactScoreQueue.size() < topK)
{
int segmentOrdinalIndex = approximateScoreQueue.pop();
int rowId = segmentOrdinalPairs.getSegmentRowId(segmentOrdinalIndex);
int ordinal = segmentOrdinalPairs.getOrdinal(segmentOrdinalIndex);
float score = reranker.similarityTo(ordinal);
exactScoreQueue.add(rowId, score);
}
rerankedCount = exactScoreQueue.size();
}
if (exactScoreQueue.size() == 0)
return endOfData();
float score = exactScoreQueue.topScore();
int rowId = exactScoreQueue.pop();
return new RowIdWithScore(rowId, score);
}
@Override
public void close()
{
FileUtils.closeQuietly(view);
}
}

125
src/java/org/apache/cassandra/index/sai/disk/v1/vector/DiskAnn.java
View File

@ -19,29 +19,25 @@
package org.apache.cassandra.index.sai.disk.v1.vector;
import java.io.IOException;
import java.util.Arrays;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.PrimitiveIterator;
import java.util.stream.IntStream;
import java.util.function.IntConsumer;
import io.github.jbellis.jvector.disk.CachingGraphIndex;
import io.github.jbellis.jvector.disk.OnDiskGraphIndex;
import io.github.jbellis.jvector.graph.GraphIndex;
import io.github.jbellis.jvector.graph.GraphSearcher;
import io.github.jbellis.jvector.graph.NeighborSimilarity;
import io.github.jbellis.jvector.graph.SearchResult;
import io.github.jbellis.jvector.graph.SearchResult.NodeScore;
import io.github.jbellis.jvector.pq.CompressedVectors;
import io.github.jbellis.jvector.util.Bits;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
import org.apache.cassandra.index.sai.disk.format.IndexComponent;
import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles;
import org.apache.cassandra.index.sai.disk.v1.postings.VectorPostingList;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.io.util.FileHandle;
import org.apache.cassandra.tracing.Tracing;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.Throwables;
public class DiskAnn implements AutoCloseable
{
@ -49,13 +45,15 @@ public class DiskAnn implements AutoCloseable
private final OnDiskOrdinalsMap ordinalsMap;
private final CachingGraphIndex graph;
private final VectorSimilarityFunction similarityFunction;
private final String source;
// only one of these will be not null
private final CompressedVectors compressedVectors;
public DiskAnn(SegmentMetadata.ComponentMetadataMap componentMetadatas, PerColumnIndexFiles indexFiles, IndexWriterConfig config) throws IOException
public DiskAnn(SegmentMetadata.ComponentMetadataMap componentMetadatas, PerColumnIndexFiles indexFiles, IndexWriterConfig config, SSTableId sstableId) throws IOException
{
similarityFunction = config.getSimilarityFunction();
source = sstableId.toString();
SegmentMetadata.ComponentMetadata termsMetadata = componentMetadatas.get(IndexComponent.TERMS_DATA);
graphHandle = indexFiles.termsData();
@ -86,86 +84,56 @@ public class DiskAnn implements AutoCloseable
return graph.size();
}
public CompressedVectors getCompressedVectors()
{
return compressedVectors;
}
/**
* @return Row IDs associated with the topK vectors near the query
*/
public VectorPostingList search(float[] queryVector, int topK, int limit, Bits acceptBits)
public CloseableIterator<RowIdWithScore> search(float[] queryVector, int topK, int limit, Bits acceptBits, IntConsumer nodesVisitedConsumer)
{
OnHeapGraph.validateIndexable(queryVector, similarityFunction);
GraphIndex.View<float[]> view = graph.getView();
GraphSearcher<float[]> searcher = new GraphSearcher.Builder<>(view).build();
NeighborSimilarity.ScoreFunction scoreFunction;
NeighborSimilarity.ReRanker<float[]> reRanker;
if (compressedVectors == null)
try
{
scoreFunction = (NeighborSimilarity.ExactScoreFunction)
i -> similarityFunction.compare(queryVector, view.getVector(i));
reRanker = null;
}
else
{
scoreFunction = compressedVectors.approximateScoreFunctionFor(queryVector, similarityFunction);
reRanker = (i, map) -> similarityFunction.compare(queryVector, map.get(i));
}
SearchResult result = searcher.search(scoreFunction,
reRanker,
topK,
ordinalsMap.ignoringDeleted(acceptBits));
Tracing.trace("DiskANN search visited {} nodes to return {} results", result.getVisitedCount(), result.getNodes().length);
return annRowIdsToPostings(result, limit);
}
private class RowIdIterator implements PrimitiveIterator.OfInt, AutoCloseable
{
private final Iterator<NodeScore> it;
private final OnDiskOrdinalsMap.RowIdsView rowIdsView = ordinalsMap.getRowIdsView();
private OfInt segmentRowIdIterator = IntStream.empty().iterator();
public RowIdIterator(NodeScore[] results)
{
this.it = Arrays.stream(results).iterator();
}
@Override
public boolean hasNext()
{
while (!segmentRowIdIterator.hasNext() && it.hasNext())
GraphSearcher<float[]> searcher = new GraphSearcher.Builder<>(view).build();
NeighborSimilarity.ScoreFunction scoreFunction;
NeighborSimilarity.ReRanker<float[]> reRanker;
if (compressedVectors == null)
{
try
{
int ordinal = it.next().node;
segmentRowIdIterator = Arrays.stream(rowIdsView.getSegmentRowIdsMatching(ordinal)).iterator();
}
catch (IOException e)
{
throw new RuntimeException(e);
}
scoreFunction = (NeighborSimilarity.ExactScoreFunction)
i -> similarityFunction.compare(queryVector, view.getVector(i));
reRanker = null;
}
return segmentRowIdIterator.hasNext();
else
{
scoreFunction = compressedVectors.approximateScoreFunctionFor(queryVector, similarityFunction);
reRanker = (i, map) -> similarityFunction.compare(queryVector, map.get(i));
}
Bits acceptedBits = ordinalsMap.ignoringDeleted(acceptBits);
// Search is done within the iterator to keep track of visited nodes. The resulting iterator
// searches until the graph is exhausted.
AutoResumingNodeScoreIterator nodeScoreIterator = new AutoResumingNodeScoreIterator(searcher, scoreFunction, reRanker, topK, acceptedBits, nodesVisitedConsumer, false, source, view);
return new NodeScoreToRowIdWithScoreIterator(nodeScoreIterator, ordinalsMap.getRowIdsView());
}
@Override
public int nextInt() {
if (!hasNext())
throw new NoSuchElementException();
return segmentRowIdIterator.nextInt();
}
@Override
public void close()
catch (Throwable e)
{
rowIdsView.close();
FileUtils.closeQuietly(view);
throw Throwables.unchecked(e);
}
}
private VectorPostingList annRowIdsToPostings(SearchResult results, int limit)
public NeighborSimilarity.ApproximateScoreFunction getApproximateScoreFunction(float[] queryVector)
{
try (var iterator = new RowIdIterator(results.getNodes()))
{
return new VectorPostingList(iterator, limit, results.getVisitedCount());
}
return compressedVectors.approximateScoreFunctionFor(queryVector, similarityFunction);
}
public NeighborSimilarity.ExactScoreFunction getExactScoreFunction(float[] queryVector, GraphIndex.View<float[]> view)
{
return i -> similarityFunction.compare(queryVector, view.getVector(i));
}
@Override
@ -180,4 +148,13 @@ public class DiskAnn implements AutoCloseable
{
return ordinalsMap.getOrdinalsView();
}
/**
* Get the graph view, callers must close the view.
* @return
*/
public GraphIndex.View<float[]> getView()
{
return graph.getView();
}
}

46
src/java/org/apache/cassandra/index/sai/disk/v1/vector/NeighborQueueRowIdIterator.java Normal file
View File

@ -0,0 +1,46 @@
/*
* 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 org.apache.cassandra.utils.AbstractIterator;
import io.github.jbellis.jvector.graph.NeighborQueue;
/**
* An iterator over {@link RowIdWithScore} that lazily consumes a {@link NeighborQueue}.
*/
public class NeighborQueueRowIdIterator extends AbstractIterator<RowIdWithScore>
{
private final NeighborQueue scoreQueue;
public NeighborQueueRowIdIterator(NeighborQueue scoreQueue)
{
this.scoreQueue = scoreQueue;
}
@Override
protected RowIdWithScore computeNext()
{
if (scoreQueue.size() == 0)
return endOfData();
float score = scoreQueue.topScore();
int rowId = scoreQueue.pop();
return new RowIdWithScore(rowId, score);
}
}

82
src/java/org/apache/cassandra/index/sai/disk/v1/vector/NodeScoreToRowIdWithScoreIterator.java Normal file
View File

@ -0,0 +1,82 @@
/*
* 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.PrimitiveIterator;
import java.util.stream.IntStream;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.utils.AbstractIterator;
import org.apache.cassandra.utils.CloseableIterator;
import io.github.jbellis.jvector.graph.SearchResult;
/**
* An iterator over {@link RowIdWithScore} sorted by score descending. The iterator converts ordinals (node ids) to
* segment row ids and pairs them with the score given by the index.
*/
public class NodeScoreToRowIdWithScoreIterator extends AbstractIterator<RowIdWithScore>
{
private final CloseableIterator<SearchResult.NodeScore> nodeScores;
private final OnDiskOrdinalsMap.RowIdsView rowIdsView;
private PrimitiveIterator.OfInt segmentRowIdIterator = IntStream.empty().iterator();
private float currentScore;
public NodeScoreToRowIdWithScoreIterator(CloseableIterator<SearchResult.NodeScore> nodeScores,
OnDiskOrdinalsMap.RowIdsView rowIdsView)
{
this.nodeScores = nodeScores;
this.rowIdsView = rowIdsView;
}
@Override
protected RowIdWithScore computeNext()
{
try
{
if (segmentRowIdIterator.hasNext())
return new RowIdWithScore(segmentRowIdIterator.nextInt(), currentScore);
while (nodeScores.hasNext())
{
SearchResult.NodeScore result = nodeScores.next();
currentScore = result.score;
int ordinal = result.node;
segmentRowIdIterator = Arrays.stream(rowIdsView.getSegmentRowIdsMatching(ordinal)).iterator();
if (segmentRowIdIterator.hasNext())
return new RowIdWithScore(segmentRowIdIterator.nextInt(), currentScore);
}
return endOfData();
}
catch (IOException e)
{
throw new RuntimeException(e);
}
}
@Override
public void close()
{
FileUtils.closeQuietly(rowIdsView);
FileUtils.closeQuietly(nodeScores);
}
}

81
src/java/org/apache/cassandra/index/sai/disk/v1/vector/OnHeapGraph.java
View File

@ -24,7 +24,6 @@ import java.util.Arrays;
import java.util.Collection;
import java.util.HashSet;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentSkipListMap;
@ -32,6 +31,7 @@ import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.IntStream;
import org.cliffc.high_scale_lib.NonBlockingHashMap;
import org.cliffc.high_scale_lib.NonBlockingHashMapLong;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@ -46,10 +46,12 @@ import io.github.jbellis.jvector.graph.SearchResult;
import io.github.jbellis.jvector.pq.CompressedVectors;
import io.github.jbellis.jvector.pq.ProductQuantization;
import io.github.jbellis.jvector.util.Bits;
import io.github.jbellis.jvector.util.RamUsageEstimator;
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.db.memtable.Memtable;
import org.apache.cassandra.exceptions.InvalidRequestException;
import org.apache.cassandra.index.sai.disk.format.IndexComponent;
import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
@ -59,45 +61,42 @@ import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
import org.apache.cassandra.io.util.SequentialWriter;
import org.apache.cassandra.tracing.Tracing;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.lucene.util.StringHelper;
public class OnHeapGraph<T>
{
private static final Logger logger = LoggerFactory.getLogger(OnHeapGraph.class);
public static final int MIN_PQ_ROWS = 1024;
private final RamAwareVectorValues vectorValues;
private final GraphIndexBuilder<float[]> builder;
private final VectorType<?> vectorType;
private final VectorSimilarityFunction similarityFunction;
private final ConcurrentMap<float[], VectorPostings<T>> postingsMap;
private final NonBlockingHashMapLong<VectorPostings<T>> postingsByOrdinal;
private final NonBlockingHashMap<T, float[]> vectorsByKey;
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);
}
private String source;
/**
* @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.
* @param memtable should be provided if attached to a memtable, null otherwise (i.e. compaction). Allows us to
* configure concurrent search and provide more meaningful trace logging. Concurrent search
* while building the graph; non-concurrent allows us to avoid synchronization costs.
*/
@SuppressWarnings("unchecked")
public OnHeapGraph(AbstractType<?> termComparator, IndexWriterConfig indexWriterConfig, boolean concurrent)
public OnHeapGraph(AbstractType<?> termComparator, IndexWriterConfig indexWriterConfig, Memtable memtable)
{
this.vectorType = (VectorType<?>) termComparator;
vectorValues = concurrent
source = memtable != null
? memtable.getClass().getSimpleName() + '@' + Integer.toHexString(memtable.hashCode())
: "compaction";
vectorValues = memtable != null
? new ConcurrentVectorValues(((VectorType<?>) termComparator).dimension)
: new CompactionVectorValues(((VectorType<Float>) termComparator));
similarityFunction = indexWriterConfig.getSimilarityFunction();
@ -107,6 +106,7 @@ public class OnHeapGraph<T>
// 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<>();
vectorsByKey = memtable != null ? new NonBlockingHashMap<>() : null;
builder = new GraphIndexBuilder<>(vectorValues,
VectorEncoding.FLOAT32,
@ -154,6 +154,16 @@ public class OnHeapGraph<T>
}
long bytesUsed = 0L;
// Store a cached reference to the vector for brute force computations later. Because insertions
// for the same primary key are guaranteed to be sequential, there is no race condition here.
if (vectorsByKey != null)
{
vectorsByKey.put(key, vector);
// The size of the entries themselves are counted below, so just count the two extra references
bytesUsed += RamUsageEstimator.NUM_BYTES_OBJECT_REF * 2L;
}
VectorPostings<T> postings = postingsMap.get(vector);
// 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:
@ -240,6 +250,13 @@ public class OnHeapGraph<T>
return postingsByOrdinal.get(node).getPostings();
}
public float[] vectorForKey(T key)
{
if (vectorsByKey == null)
throw new IllegalStateException("vectorsByKey is not initialized");
return vectorsByKey.get(key);
}
public long remove(ByteBuffer term, T key)
{
assert term != null && term.remaining() != 0;
@ -255,31 +272,35 @@ public class OnHeapGraph<T>
}
hasDeletions = true;
return postings.remove(key);
long bytesUsed = postings.remove(key);
if (vectorsByKey != null)
{
// On updates to a row, we call add then remove, so we must pass the key's value to ensure we only remove
// the deleted vector from vectorsByKey.
vectorsByKey.remove(key, vector);
}
return bytesUsed;
}
/**
* @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)
public CloseableIterator<SearchResult.NodeScore> 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<>();
return CloseableIterator.empty();
Bits bits = hasDeletions ? BitsUtil.bitsIgnoringDeleted(toAccept, postingsByOrdinal) : toAccept;
GraphIndex<float[]> graph = builder.getGraph();
GraphSearcher<float[]> searcher = new GraphSearcher.Builder<>(graph.getView()).withConcurrentUpdates().build();
GraphIndex.View<float[]> view = graph.getView();
GraphSearcher<float[]> searcher = new GraphSearcher.Builder<>(view).withConcurrentUpdates().build();
NeighborSimilarity.ExactScoreFunction scoreFunction = node2 -> vectorCompareFunction(queryVector, node2);
SearchResult result = searcher.search(scoreFunction, null, limit, bits);
Tracing.trace("ANN search visited {} in-memory nodes to return {} results", result.getVisitedCount(), result.getNodes().length);
SearchResult.NodeScore[] a = result.getNodes();
PriorityQueue<T> keyQueue = new PriorityQueue<>();
for (int i = 0; i < a.length; i++)
keyQueue.addAll(keysFromOrdinal(a[i].node));
return keyQueue;
return new AutoResumingNodeScoreIterator(searcher, scoreFunction, null, limit, bits, v -> {}, true, source, view);
}
public SegmentMetadata.ComponentMetadataMap writeData(IndexDescriptor indexDescriptor, IndexIdentifier indexIdentifier, Function<T, Integer> postingTransformer) throws IOException
@ -351,8 +372,8 @@ public class OnHeapGraph<T>
{
// 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)
writer.writeBoolean(vectorValues.size() >= MIN_PQ_ROWS);
if (vectorValues.size() < MIN_PQ_ROWS)
{
logger.debug("Skipping PQ for only {} vectors", vectorValues.size());
return writer.position();

2
src/java/org/apache/cassandra/index/sai/disk/v1/vector/OptimizeFor.java
View File

@ -42,7 +42,7 @@ public enum OptimizeFor
public int topKFor(int limit)
{
return (int)(limitMultiplier.apply(limit) * limit);
return (int)(Math.max(1.0, limitMultiplier.apply(limit)) * limit);
}
public static OptimizeFor fromString(String value)

81
src/java/org/apache/cassandra/index/sai/disk/v1/vector/PrimaryKeyWithScore.java Normal file
View File

@ -0,0 +1,81 @@
/*
* 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 org.apache.cassandra.db.CellSourceIdentifier;
import org.apache.cassandra.db.rows.Cell;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.index.sai.utils.CellWithSource;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.schema.ColumnMetadata;
/**
* A PrimaryKey with one piece of metadata. Subclasses define the metadata, and to prevent unnecessary boxing, the
* metadata is not referenced in this class. The metadata is not used to determine equality or hash code, but it is used
* to compare the PrimaryKey objects.
* Note: this class has a natural ordering that is inconsistent with equals.
*/
public class PrimaryKeyWithScore implements Comparable<PrimaryKeyWithScore>
{
protected final ColumnMetadata columnMetadata;
private final PrimaryKey primaryKey;
private final CellSourceIdentifier sourceTable;
private final float indexScore;
public PrimaryKeyWithScore(ColumnMetadata columnMetadata, CellSourceIdentifier sourceTable, PrimaryKey primaryKey, float indexScore)
{
this.columnMetadata = columnMetadata;
this.sourceTable = sourceTable;
this.primaryKey = primaryKey;
this.indexScore = indexScore;
}
public PrimaryKey primaryKey()
{
return primaryKey;
}
public boolean isIndexDataValid(Row row, long nowInSecs)
{
// If the indexed column is part of the primary key, we don't need this type of validation because we would have
// fetched the row using the indexed primary key, so they have to match.
if (columnMetadata.isPrimaryKeyColumn())
return true;
// If the row is static and the column is not static, or vice versa, the indexed value won't be present so we
// don't need to check if live data matches indexed data.
if (row.isStatic() != columnMetadata.isStatic())
return true;
Cell<?> cell = row.getCell(columnMetadata);
if (!cell.isLive(nowInSecs))
return false;
assert cell instanceof CellWithSource : "Expected CellWithSource, got " + cell.getClass();
return sourceTable.isEqualSource(((CellWithSource<?>) cell).sourceTable());
}
@Override
public int compareTo(PrimaryKeyWithScore o)
{
// Descending order
return Float.compare(o.indexScore, indexScore);
}
}

69
src/java/org/apache/cassandra/index/sai/disk/v1/vector/RowIdToPrimaryKeyWithScoreIterator.java Normal file
View File

@ -0,0 +1,69 @@
/*
* 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 org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.utils.AbstractIterator;
import org.apache.cassandra.utils.CloseableIterator;
/**
* An iterator over scored primary keys ordered by the score descending
* Not skippable.
*/
public class RowIdToPrimaryKeyWithScoreIterator extends AbstractIterator<PrimaryKeyWithScore>
{
private final ColumnMetadata column;
private final SSTableId sstableId;
private final PrimaryKeyMap primaryKeyMap;
private final CloseableIterator<RowIdWithScore> scoredRowIdIterator;
private final long segmentRowIdOffset;
public RowIdToPrimaryKeyWithScoreIterator(ColumnMetadata column,
PrimaryKeyMap.Factory primaryKeyMapFactory,
CloseableIterator<RowIdWithScore> scoredRowIdIterator,
long segmentRowIdOffset) throws IOException
{
this.column = column;
this.scoredRowIdIterator = scoredRowIdIterator;
this.primaryKeyMap = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap();
this.sstableId = primaryKeyMap.getSSTableId();
this.segmentRowIdOffset = segmentRowIdOffset;
}
@Override
protected PrimaryKeyWithScore computeNext()
{
if (!scoredRowIdIterator.hasNext())
return endOfData();
RowIdWithScore rowIdWithScore = scoredRowIdIterator.next();
return rowIdWithScore.toPrimaryKeyWithScore(column, sstableId, primaryKeyMap, segmentRowIdOffset);
}
@Override
public void close()
{
FileUtils.closeQuietly(primaryKeyMap);
FileUtils.closeQuietly(scoredRowIdIterator);
}
}

48
src/java/org/apache/cassandra/index/sai/disk/v1/vector/RowIdWithScore.java Normal file
View File

@ -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.vector;
import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.schema.ColumnMetadata;
/**
* Represents a row id with its computed score.
*/
public class RowIdWithScore
{
private final int segmentRowId;
private final float score;
public RowIdWithScore(int segmentRowId, float score)
{
this.segmentRowId = segmentRowId;
this.score = score;
}
public PrimaryKeyWithScore toPrimaryKeyWithScore(ColumnMetadata columnMetadata,
SSTableId sstableId,
PrimaryKeyMap primaryKeyMap,
long segmentRowIdOffset)
{
PrimaryKey pk = primaryKeyMap.primaryKeyFromRowId(segmentRowIdOffset + segmentRowId);
return new PrimaryKeyWithScore(columnMetadata, sstableId, pk, score);
}
}

131
src/java/org/apache/cassandra/index/sai/disk/v1/vector/SegmentRowIdOrdinalPairs.java Normal file
View File

@ -0,0 +1,131 @@
/*
* 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.util.function.IntConsumer;
import io.github.jbellis.jvector.graph.NeighborQueue;
import io.github.jbellis.jvector.graph.NeighborSimilarity;
/**
* A specialized data structure that stores segment row id to ordinal pairs efficiently. Implemented as an array of int
* pairs that avoids boxing.
*/
public class SegmentRowIdOrdinalPairs
{
private final int capacity;
private int size;
private final int[] array;
/**
* Create a new SegmentRowIdOrdinalPairs with the given capacity.
* @param capacity the capacity
*/
public SegmentRowIdOrdinalPairs(int capacity)
{
assert capacity < Integer.MAX_VALUE / 2 : "capacity is too large " + capacity;
this.capacity = capacity;
this.size = 0;
this.array = new int[capacity * 2];
}
/**
* Add a pair to the array.
* @param segmentRowId the first value
* @param ordinal the second value
*/
public void add(int segmentRowId, int ordinal)
{
if (size == capacity)
throw new ArrayIndexOutOfBoundsException(size);
array[size * 2] = segmentRowId;
array[size * 2 + 1] = ordinal;
size++;
}
/**
* Get the row id at the given index.
* @param index the index
* @return the row id
*/
public int getSegmentRowId(int index)
{
if ( index < 0 || index >= size)
throw new ArrayIndexOutOfBoundsException(index);
return array[index * 2];
}
/**
* Get the ordinal at the given index.
* @param index the index
* @return the ordinal
*/
public int getOrdinal(int index)
{
if ( index < 0 || index >= size)
throw new ArrayIndexOutOfBoundsException(index);
return array[index * 2 + 1];
}
/**
* The number of pairs in the array.
* @return the number of pairs in the array
*/
public int size()
{
return size;
}
/**
* Create an iterator over the segment row id and scored ordinal pairs in the array.
* @param scoreFunction the score function to use to compute the next score based on the ordinal
* @return a {@link NeighborQueue}
*/
public NeighborQueue mapToSegmentRowIdScoreHeap(NeighborSimilarity.ScoreFunction scoreFunction)
{
// TODO this could be improved using Floyd's algorithm in a later jvector version
NeighborQueue queue = new NeighborQueue(size(), true);
for (int i = 0; i < size; i++)
queue.add(array[i * 2], scoreFunction.similarityTo(array[i * 2 + 1])); // rowid, score
return queue;
}
/**
* Create an iterator over the index and scored ordinal pairs in the array.
* @param scoreFunction the score function to use to compute the next score based on the ordinal
*/
public NeighborQueue mapToIndexScoreIterator(NeighborSimilarity.ScoreFunction scoreFunction)
{
// TODO this could be improved using Floyd's algorithm in a later jvector version
NeighborQueue queue = new NeighborQueue(size(), true);
for (int i = 0; i < size; i++)
queue.add(i, scoreFunction.similarityTo(array[i * 2 + 1])); // index, score
return queue;
}
/**
* Calls the consumer for each right value in each pair of the array.
* @param consumer the consumer to call for each right value
*/
public void forEachOrdinal(IntConsumer consumer)
{
for (int i = 0; i < size; i++)
consumer.accept(array[i * 2 + 1]);
}
}

67
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeListIterator.java
View File

@ -1,67 +0,0 @@
/*
* 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.
*
* @param minimumKey the minimum key for the provided list of keys
* @param maximumKey the maximum key for the provided list of keys
* @param keys the list of keys to iterate over
*/
public KeyRangeListIterator(PrimaryKey minimumKey, PrimaryKey maximumKey, List<PrimaryKey> keys)
{
super(minimumKey, maximumKey, keys.size());
this.keyQueue = Iterators.peekingIterator(keys.iterator());
}
@Override
protected void performSkipTo(PrimaryKey nextKey)
{
while (keyQueue.hasNext())
{
if (keyQueue.peek().compareTo(nextKey, false) >= 0)
break;
keyQueue.next();
}
}
@Override
public void close() {}
@Override
protected PrimaryKey computeNext()
{
return keyQueue.hasNext() ? keyQueue.next() : endOfData();
}
}

93
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeOrderingIterator.java
View File

@ -1,93 +0,0 @@
/*
* 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.ArrayList;
import java.util.List;
import java.util.function.Function;
import javax.annotation.concurrent.NotThreadSafe;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.util.FileUtils;
/**
* An iterator that consumes a chunk of {@link PrimaryKey}s from the {@link KeyRangeIterator}, passes them to the
* {@link Function} to filter the chunk of {@link PrimaryKey}s and then pass the results to next consumer.
* The PKs are currently returned in {@link PrimaryKey} order, but that contract may change.
*/
@NotThreadSafe
public class KeyRangeOrderingIterator extends KeyRangeIterator
{
private final KeyRangeIterator input;
private final int chunkSize;
private final Function<List<PrimaryKey>, KeyRangeIterator> nextRangeFunction;
private final ArrayList<PrimaryKey> nextKeys;
private KeyRangeIterator nextIterator;
public KeyRangeOrderingIterator(KeyRangeIterator input, int chunkSize, Function<List<PrimaryKey>, KeyRangeIterator> nextRangeFunction)
{
super(input, () -> {});
this.input = input;
this.chunkSize = chunkSize;
this.nextRangeFunction = nextRangeFunction;
this.nextKeys = new ArrayList<>(chunkSize);
}
@Override
public PrimaryKey computeNext()
{
if (nextIterator == null || !nextIterator.hasNext())
{
do
{
if (!input.hasNext())
return endOfData();
nextKeys.clear();
do
{
nextKeys.add(input.next());
}
while (nextKeys.size() < chunkSize && input.hasNext());
// Get the next iterator before closing this one to prevent releasing the resource.
KeyRangeIterator previousIterator = nextIterator;
// If this results in an exception, previousIterator is closed in close() method.
nextIterator = nextRangeFunction.apply(nextKeys);
if (previousIterator != null)
FileUtils.closeQuietly(previousIterator);
// nextIterator might not have any rows due to shadowed primary keys
}
while (!nextIterator.hasNext());
}
return nextIterator.next();
}
@Override
protected void performSkipTo(PrimaryKey nextToken)
{
input.skipTo(nextToken);
if (nextIterator != null)
nextIterator.skipTo(nextToken);
}
public void close()
{
FileUtils.closeQuietly(input);
FileUtils.closeQuietly(nextIterator);
}
}

47
src/java/org/apache/cassandra/index/sai/iterators/PriorityQueueIterator.java Normal file
View File

@ -0,0 +1,47 @@
/*
* 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.PriorityQueue;
import org.apache.cassandra.utils.AbstractIterator;
/**
* An iterator over a priority queue.
* @param <T> the type of the elements in the priority queue
*/
public class PriorityQueueIterator<T> extends AbstractIterator<T>
{
private final PriorityQueue<T> queue;
/**
* Build a PriorityQueueIterator.
* @param queue a priority queue to be lazily consumed by the iterator
*/
public PriorityQueueIterator(PriorityQueue<T> queue)
{
this.queue = queue;
}
@Override
protected T computeNext()
{
return queue.isEmpty() ? endOfData() : queue.poll();
}
}

27
src/java/org/apache/cassandra/index/sai/memory/MemtableIndex.java
View File

@ -28,17 +28,19 @@ import java.util.function.Function;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.marshal.AbstractType;
import org.apache.cassandra.db.memtable.Memtable;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
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.utils.CloseableIterator;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;
@ -47,12 +49,12 @@ public class MemtableIndex implements MemtableOrdering
private final MemoryIndex memoryIndex;
private final LongAdder writeCount = new LongAdder();
private final LongAdder estimatedMemoryUsed = new LongAdder();
private final AbstractType<?> type;
private final Memtable memtable;
public MemtableIndex(StorageAttachedIndex index)
public MemtableIndex(StorageAttachedIndex index, Memtable memtable)
{
this.memoryIndex = index.termType().isVector() ? new VectorMemoryIndex(index) : new TrieMemoryIndex(index);
this.type = index.termType().indexType();
this.memoryIndex = index.termType().isVector() ? new VectorMemoryIndex(index, memtable) : new TrieMemoryIndex(index);
this.memtable = memtable;
}
public long writeCount()
@ -70,6 +72,11 @@ public class MemtableIndex implements MemtableOrdering
return memoryIndex.isEmpty();
}
public Memtable getMemtable()
{
return memtable;
}
public ByteBuffer getMinTerm()
{
return memoryIndex.getMinTerm();
@ -114,8 +121,14 @@ public class MemtableIndex implements MemtableOrdering
}
@Override
public KeyRangeIterator limitToTopResults(List<PrimaryKey> primaryKeys, Expression expression, int limit)
public CloseableIterator<PrimaryKeyWithScore> orderBy(QueryContext queryContext, Expression orderer, AbstractBounds<PartitionPosition> keyRange)
{
return memoryIndex.limitToTopResults(primaryKeys, expression, limit);
return memoryIndex.orderBy(queryContext, orderer, keyRange);
}
@Override
public CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext queryContext, List<PrimaryKey> results, Expression orderer)
{
return memoryIndex.orderResultsBy(queryContext, results, orderer);
}
}

22
src/java/org/apache/cassandra/index/sai/memory/MemtableIndexManager.java
View File

@ -50,15 +50,25 @@ public class MemtableIndexManager
this.liveMemtableIndexMap = new ConcurrentHashMap<>();
}
public long index(DecoratedKey key, Row row, Memtable mt)
public void maybeInitializeMemtableIndex(Memtable memtable)
{
if (index.termType().isVector())
initializeMemtableIndex(memtable);
}
private MemtableIndex initializeMemtableIndex(Memtable mt)
{
MemtableIndex current = liveMemtableIndexMap.get(mt);
// We expect the relevant IndexMemtable to be present most of the time, so only make the
// call to computeIfAbsent() if it's not. (see https://bugs.openjdk.java.net/browse/JDK-8161372)
MemtableIndex target = (current != null)
? current
: liveMemtableIndexMap.computeIfAbsent(mt, memtable -> new MemtableIndex(index));
return current != null ? current
: liveMemtableIndexMap.computeIfAbsent(mt, memtable -> new MemtableIndex(index, memtable));
}
public long index(DecoratedKey key, Row row, Memtable mt)
{
MemtableIndex target = initializeMemtableIndex(mt);
long start = Clock.Global.nanoTime();
@ -92,9 +102,9 @@ public class MemtableIndexManager
return index(key, newRow, memtable);
}
// Updates should only be able to happen on memtables that were already created and that are still live.
MemtableIndex target = liveMemtableIndexMap.get(memtable);
if (target == null)
return 0;
assert target != null : "Memtable for " + memtable.metadata().getTableName() + " not found";
ByteBuffer oldValue = index.termType().valueOf(key, oldRow, FBUtilities.nowInSeconds());
ByteBuffer newValue = index.termType().valueOf(key, newRow, FBUtilities.nowInSeconds());

31
src/java/org/apache/cassandra/index/sai/memory/MemtableOrdering.java
View File

@ -20,9 +20,13 @@ package org.apache.cassandra.index.sai.memory;
import java.util.List;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.utils.CloseableIterator;
/**
* Analogue of {@link org.apache.cassandra.index.sai.disk.v1.segment.SegmentOrdering}, but for memtables.
@ -30,13 +34,22 @@ import org.apache.cassandra.index.sai.utils.PrimaryKey;
public interface MemtableOrdering
{
/**
* Filter the given list of {@code PrimaryKey} results to the top `limit` results corresponding to the given expression,
* Returns an iterator over the results that is put back in token order.
* <p>
* Assumes that the given list spans the same rows as the implementing index's segment.
* Order the index based on the given orderer (expression).
*
* @param queryContext - the query context
* @param orderer - the expression to order by
* @param keyRange - the key range to search
* @return an iterator over the results in score order.
*/
default KeyRangeIterator limitToTopResults(List<PrimaryKey> primaryKeys, Expression expression, int limit)
{
throw new UnsupportedOperationException();
}
CloseableIterator<PrimaryKeyWithScore> orderBy(QueryContext queryContext,
Expression orderer,
AbstractBounds<PartitionPosition> keyRange);
/**
* Order the given list of {@link PrimaryKey} results corresponding to the given orderer.
* Returns an iterator over the results in score order.
*
* Assumes that the given spans the same rows as the implementing index's segment.
*/
CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer);
}

15
src/java/org/apache/cassandra/index/sai/memory/TrieMemoryIndex.java
View File

@ -20,6 +20,7 @@ package org.apache.cassandra.index.sai.memory;
import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.SortedSet;
@ -42,11 +43,13 @@ import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.analyzer.AbstractAnalyzer;
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.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.PrimaryKeys;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;
@ -260,6 +263,18 @@ public class TrieMemoryIndex extends MemoryIndex
: new FilteringInMemoryKeyRangeIterator(primaryKeys.keys(), keyRange);
}
@Override
public CloseableIterator<PrimaryKeyWithScore> orderBy(QueryContext queryContext, Expression orderer, AbstractBounds<PartitionPosition> keyRange)
{
throw new UnsupportedOperationException();
}
@Override
public CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer)
{
throw new UnsupportedOperationException();
}
private static class Collector
{
final PriorityQueue<PrimaryKey> mergedKeys;

201
src/java/org/apache/cassandra/index/sai/memory/VectorMemoryIndex.java
View File

@ -20,7 +20,9 @@ package org.apache.cassandra.index.sai.memory;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.NavigableSet;
@ -32,25 +34,31 @@ import java.util.function.Function;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import io.github.jbellis.jvector.graph.SearchResult;
import io.github.jbellis.jvector.util.Bits;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.memtable.Memtable;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.VectorQueryContext;
import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
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.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.iterators.KeyRangeListIterator;
import org.apache.cassandra.index.sai.iterators.PriorityQueueIterator;
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.io.util.FileUtils;
import org.apache.cassandra.tracing.Tracing;
import org.apache.cassandra.utils.AbstractIterator;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;
@ -62,6 +70,7 @@ import static java.lang.Math.pow;
public class VectorMemoryIndex extends MemoryIndex
{
private final OnHeapGraph<PrimaryKey> graph;
private final Memtable memtable;
private final LongAdder writeCount = new LongAdder();
private PrimaryKey minimumKey;
@ -69,10 +78,11 @@ public class VectorMemoryIndex extends MemoryIndex
private final NavigableSet<PrimaryKey> primaryKeys = new ConcurrentSkipListSet<>();
public VectorMemoryIndex(StorageAttachedIndex index)
public VectorMemoryIndex(StorageAttachedIndex index, Memtable memtable)
{
super(index);
this.graph = new OnHeapGraph<>(index.termType().indexType(), index.indexWriterConfig());
this.graph = new OnHeapGraph<>(index.termType().indexType(), index.indexWriterConfig(), memtable);
this.memtable = memtable;
}
@Override
@ -149,12 +159,16 @@ public class VectorMemoryIndex extends MemoryIndex
}
@Override
public KeyRangeIterator search(QueryContext queryContext, Expression expr, AbstractBounds<PartitionPosition> keyRange)
public KeyRangeIterator search(QueryContext queryContext, Expression expression, AbstractBounds<PartitionPosition> keyRange)
{
throw new UnsupportedOperationException();
}
@Override
public CloseableIterator<PrimaryKeyWithScore> orderBy(QueryContext queryContext, Expression expr, AbstractBounds<PartitionPosition> keyRange)
{
assert expr.getIndexOperator() == Expression.IndexOperator.ANN : "Only ANN is supported for vector search, received " + expr.getIndexOperator();
VectorQueryContext vectorQueryContext = queryContext.vectorContext();
ByteBuffer buffer = expr.lower().value.raw;
float[] qv = index.termType().decomposeVector(buffer);
@ -172,73 +186,93 @@ public class VectorMemoryIndex extends MemoryIndex
PrimaryKey right = isMaxToken ? null : index.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();
return CloseableIterator.empty();
int bruteForceRows = maxBruteForceRows(vectorQueryContext.limit(), resultKeys.size(), graph.size());
int bruteForceRows = maxBruteForceRows(queryContext.limit(), resultKeys.size(), graph.size());
Tracing.trace("Search range covers {} rows; max brute force rows is {} for memtable index with {} nodes, LIMIT {}",
resultKeys.size(), bruteForceRows, graph.size(), vectorQueryContext.limit());
if (resultKeys.size() < Math.max(vectorQueryContext.limit(), bruteForceRows))
return new ReorderingRangeIterator(new PriorityQueue<>(resultKeys));
resultKeys.size(), bruteForceRows, graph.size(), queryContext.limit());
if (resultKeys.size() < Math.max(queryContext.limit(), bruteForceRows))
return orderByBruteForce(qv, resultKeys);
else
bits = new KeyRangeFilteringBits(keyRange, vectorQueryContext.bitsetForShadowedPrimaryKeys(graph));
bits = new KeyRangeFilteringBits(keyRange, null);
}
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);
// Accept all bits
bits = new Bits.MatchAllBits(Integer.MAX_VALUE);
}
PriorityQueue<PrimaryKey> keyQueue = graph.search(qv, queryContext.vectorContext().limit(), bits);
if (keyQueue.isEmpty())
return KeyRangeIterator.empty();
return new ReorderingRangeIterator(keyQueue);
CloseableIterator<SearchResult.NodeScore> iterator = graph.search(qv, queryContext.limit(), bits);
return new NodeScoreToScoredPrimaryKeyIterator(iterator);
}
@Override
public KeyRangeIterator limitToTopResults(List<PrimaryKey> primaryKeys, Expression expression, int limit)
public CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext queryContext, List<PrimaryKey> results, Expression orderer)
{
if (minimumKey == null)
// This case implies maximumKey is empty too.
return KeyRangeIterator.empty();
return CloseableIterator.empty();
List<PrimaryKey> results = primaryKeys.stream()
.dropWhile(k -> k.compareTo(minimumKey) < 0)
.takeWhile(k -> k.compareTo(maximumKey) <= 0)
.collect(Collectors.toList());
int limit = queryContext.limit();
int maxBruteForceRows = maxBruteForceRows(limit, results.size(), graph.size());
List<PrimaryKey> resultsInRange = results.stream()
.dropWhile(k -> k.compareTo(minimumKey) < 0)
.takeWhile(k -> k.compareTo(maximumKey) <= 0)
.collect(Collectors.toList());
int maxBruteForceRows = maxBruteForceRows(limit, resultsInRange.size(), graph.size());
Tracing.trace("SAI materialized {} rows; max brute force rows is {} for memtable index with {} nodes, LIMIT {}",
results.size(), maxBruteForceRows, graph.size(), limit);
if (results.size() <= maxBruteForceRows)
{
if (results.isEmpty())
return KeyRangeIterator.empty();
return new KeyRangeListIterator(minimumKey, maximumKey, results);
}
resultsInRange.size(), maxBruteForceRows, graph.size(), limit);
ByteBuffer buffer = expression.lower().value.raw;
if (resultsInRange.isEmpty())
return CloseableIterator.empty();
ByteBuffer buffer = orderer.lower().value.raw;
float[] qv = index.termType().decomposeVector(buffer);
KeyFilteringBits bits = new KeyFilteringBits(results);
PriorityQueue<PrimaryKey> keyQueue = graph.search(qv, limit, bits);
if (keyQueue.isEmpty())
return KeyRangeIterator.empty();
return new ReorderingRangeIterator(keyQueue);
if (resultsInRange.size() <= maxBruteForceRows)
return orderByBruteForce(qv, resultsInRange);
// Search the graph for the topK vectors near the query
KeyFilteringBits bits = new KeyFilteringBits(resultsInRange);
CloseableIterator<SearchResult.NodeScore> nodeScores = graph.search(qv, limit, bits);
return new NodeScoreToScoredPrimaryKeyIterator(nodeScores);
}
private int maxBruteForceRows(int limit, int nPermittedOrdinals, int graphSize)
{
int expectedNodesVisited = expectedNodesVisited(limit, nPermittedOrdinals, graphSize);
int expectedComparisons = index.indexWriterConfig().getMaximumNodeConnections() * expectedNodesVisited;
// in-memory comparisons are cheaper than pulling a row off disk and then comparing
// VSTODO this is dramatically oversimplified
// larger dimension should increase this, because comparisons are more expensive
// lower chunk cache hit ratio should decrease this, because loading rows is more expensive
double memoryToDiskFactor = 0.25;
return (int) max(limit, memoryToDiskFactor * expectedComparisons);
// ANN index will do a bunch of extra work besides the full comparisons
// VSTODO I'm not sure which one is more expensive, but since the graph is in memory and the vectors are
// full precision, the goal here is simple: minimize the number of vector comparisons (aka nodes visited).
// As such, the cost function weights them at a 1:1 ratio for now.
return max(limit, expectedNodesVisited);
}
private CloseableIterator<PrimaryKeyWithScore> orderByBruteForce(float[] queryVector, Collection<PrimaryKey> keys)
{
VectorSimilarityFunction similarityFunction = index.indexWriterConfig().getSimilarityFunction();
List<PrimaryKeyWithScore> scoredKeys = new ArrayList<>(keys.size());
for (PrimaryKey key : keys)
{
PrimaryKeyWithScore scoredKey = scoreKey(similarityFunction, queryVector, key);
if (scoredKey != null)
scoredKeys.add(scoredKey);
}
// Because we merge iterators from all sstables and memtables, we do not need a complete sort of these
// elements, so a priority queue provides good performance.
return new PriorityQueueIterator<>(new PriorityQueue<>(scoredKeys));
}
private PrimaryKeyWithScore scoreKey(VectorSimilarityFunction similarityFunction, float[] queryVector, PrimaryKey key)
{
float[] vector = graph.vectorForKey(key);
if (vector == null)
return null;
float score = similarityFunction.compare(queryVector, vector);
return new PrimaryKeyWithScore(index.termType().columnMetadata(), memtable, key, score);
}
/**
@ -321,36 +355,6 @@ public class VectorMemoryIndex extends MemoryIndex
}
}
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)
{
while (!keyQueue.isEmpty() && keyQueue.peek().compareTo(nextKey) < 0)
keyQueue.poll();
}
@Override
public void close() {}
@Override
protected PrimaryKey computeNext()
{
if (keyQueue.isEmpty())
return endOfData();
return keyQueue.poll();
}
}
private class KeyFilteringBits implements Bits
{
private final List<PrimaryKey> results;
@ -373,4 +377,45 @@ public class VectorMemoryIndex extends MemoryIndex
return results.size();
}
}
/**
* An iterator over {@link PrimaryKeyWithScore} sorted by score descending. The iterator converts ordinals (node ids)
* to {@link PrimaryKey}s and pairs them with the score given by the index.
*/
private class NodeScoreToScoredPrimaryKeyIterator extends AbstractIterator<PrimaryKeyWithScore>
{
private final CloseableIterator<SearchResult.NodeScore> nodeScores;
private Iterator<PrimaryKeyWithScore> primaryKeysForNode = Collections.emptyIterator();
NodeScoreToScoredPrimaryKeyIterator(CloseableIterator<SearchResult.NodeScore> nodeScores)
{
this.nodeScores = nodeScores;
}
@Override
protected PrimaryKeyWithScore computeNext()
{
if (primaryKeysForNode.hasNext())
return primaryKeysForNode.next();
while (nodeScores.hasNext())
{
SearchResult.NodeScore nodeScore = nodeScores.next();
primaryKeysForNode = graph.keysFromOrdinal(nodeScore.node)
.stream()
.map(pk -> new PrimaryKeyWithScore(index.termType().columnMetadata(), memtable, pk, nodeScore.score))
.iterator();
if (primaryKeysForNode.hasNext())
return primaryKeysForNode.next();
}
return endOfData();
}
@Override
public void close()
{
FileUtils.closeQuietly(nodeScores);
}
}
}

28
src/java/org/apache/cassandra/index/sai/plan/Operation.java
View File

@ -26,7 +26,6 @@ import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ArrayListMultimap;
@ -42,9 +41,11 @@ import org.apache.cassandra.exceptions.InvalidRequestException;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.analyzer.AbstractAnalyzer;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.utils.IndexTermType;
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.utils.CloseableIterator;
public class Operation
{
@ -319,16 +320,23 @@ public class Operation
*/
static KeyRangeIterator buildIterator(QueryController controller)
{
List<RowFilter.Expression> orderings = controller.indexFilter().getExpressions()
.stream().filter(e -> e.operator() == Operator.ANN).collect(Collectors.toList());
assert orderings.size() <= 1;
if (controller.indexFilter().getExpressions().size() == 1 && orderings.size() == 1)
return Node.buildTree(controller.indexFilter()).analyzeTree(controller).rangeIterator(controller);
}
/**
* Converts expressions into filter tree for query.
*
* @return a KeyRangeIterator over the index query results
*/
static CloseableIterator<PrimaryKeyWithScore> buildIteratorForOrder(QueryController controller, QueryViewBuilder.QueryExpressionView view)
{
if (controller.indexFilter().getExpressions().size() == 1)
// If we only have one expression, we just use the ANN index to order and limit.
return controller.getTopKRows(orderings.get(0));
KeyRangeIterator iterator = Node.buildTree(controller.indexFilter()).analyzeTree(controller).rangeIterator(controller);
if (orderings.isEmpty())
return iterator;
return controller.getTopKRows(iterator, orderings.get(0));
return controller.getTopKRows(view);
// Otherwise, we need to search first, then order.
KeyRangeIterator iterator = buildIterator(controller);
return controller.getTopKRows(iterator, view);
}
/**

224
src/java/org/apache/cassandra/index/sai/plan/QueryController.java
View File

@ -18,20 +18,20 @@
package org.apache.cassandra.index.sai.plan;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.NavigableSet;
import java.util.function.Function;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import com.google.common.collect.Lists;
import org.apache.cassandra.cql3.Operator;
import org.apache.cassandra.config.CassandraRelevantProperties;
import org.apache.cassandra.db.CellSourceIdentifier;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.DataRange;
@ -47,30 +47,54 @@ import org.apache.cassandra.db.filter.ClusteringIndexNamesFilter;
import org.apache.cassandra.db.filter.DataLimits;
import org.apache.cassandra.db.filter.RowFilter;
import org.apache.cassandra.db.guardrails.Guardrails;
import org.apache.cassandra.db.rows.BaseRowIterator;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.db.rows.UnfilteredRowIterator;
import org.apache.cassandra.db.transform.Transformation;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.VectorQueryContext;
import org.apache.cassandra.index.sai.disk.IndexSearchResultIterator;
import org.apache.cassandra.index.sai.disk.SSTableIndex;
import org.apache.cassandra.index.sai.iterators.KeyRangeConcatIterator;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.iterators.KeyRangeOrderingIterator;
import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator;
import org.apache.cassandra.index.sai.memory.MemtableIndex;
import org.apache.cassandra.index.sai.utils.MergePrimaryKeyWithScoreIterator;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.RowWithSource;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.net.ParamType;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.tracing.Tracing;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.InsertionOrderedNavigableSet;
import org.apache.cassandra.utils.Throwables;
import static org.apache.cassandra.config.CassandraRelevantProperties.SAI_VECTOR_SEARCH_ORDER_CHUNK_SIZE;
public class QueryController
{
// Transforms a row to include its source table, which is then used for ANN query validation.
private final static Function<CellSourceIdentifier, Transformation<BaseRowIterator<?>>> SOURCE_TABLE_ROW_TRANSFORMER = (CellSourceIdentifier sourceTable) -> new Transformation<>()
{
@Override
protected Row applyToStatic(Row row)
{
return new RowWithSource(row, sourceTable);
}
@Override
protected Row applyToRow(Row row)
{
return new RowWithSource(row, sourceTable);
}
};
/**
* The maximum number of primary keys we will materialize when performing hybrid vector search. If this limit is
* exceeded, we switch to an order-by-then-filter execution path
*/
public static int MAX_MATERIALIZED_KEYS = CassandraRelevantProperties.SAI_VECTOR_SEARCH_MAX_MATERIALIZE_KEYS.getInt();
final QueryContext queryContext;
private final ColumnFamilyStore cfs;
@ -81,7 +105,6 @@ public class QueryController
private final PrimaryKey.Factory keyFactory;
private final PrimaryKey firstPrimaryKey;
private final PrimaryKey lastPrimaryKey;
private final int orderChunkSize;
private final NavigableSet<Clustering<?>> nextClusterings;
@ -101,7 +124,6 @@ public class QueryController
this.keyFactory = new PrimaryKey.Factory(cfs.getPartitioner(), cfs.getComparator());
this.firstPrimaryKey = keyFactory.create(mergeRange.left.getToken());
this.lastPrimaryKey = keyFactory.create(mergeRange.right.getToken());
this.orderChunkSize = SAI_VECTOR_SEARCH_ORDER_CHUNK_SIZE.getInt();
this.nextClusterings = new InsertionOrderedNavigableSet<>(cfs.metadata().comparator);
}
@ -143,6 +165,11 @@ public class QueryController
return ranges;
}
public AbstractBounds<PartitionPosition> mergeRange()
{
return mergeRange;
}
@Nullable
public StorageAttachedIndex indexFor(RowFilter.Expression expression)
{
@ -171,6 +198,31 @@ public class QueryController
return partition.queryMemtableAndDisk(cfs, executionController);
}
/**
* Get an iterator over the row(s) for this primary key. Restrict the search to the specified view. Apply the
* {@link #SOURCE_TABLE_ROW_TRANSFORMER} so that resulting cells have the source memtable/sstable. Expect one row
* for a fully qualified primary key or all rows within a partition for a static primary key.
*
* @param key primary key to fetch from storage.
* @param executionController the executionController to use when querying storage
* @return an iterator of rows matching the query
*/
public UnfilteredRowIterator queryStorage(PrimaryKey key, ColumnFamilyStore.ViewFragment view, ReadExecutionController executionController)
{
if (key == null)
throw new IllegalArgumentException("non-null key required");
SinglePartitionReadCommand partition = SinglePartitionReadCommand.create(cfs.metadata(),
command.nowInSec(),
command.columnFilter(),
RowFilter.none(),
DataLimits.NONE,
key.partitionKey(),
makeFilter(List.of(key)));
return partition.queryMemtableAndDisk(cfs, view, SOURCE_TABLE_ROW_TRANSFORMER, executionController);
}
/**
* Build a {@link KeyRangeIterator.Builder} from the given list of {@link Expression}s.
* <p>
@ -199,10 +251,9 @@ public class QueryController
expressions = expressions.stream().filter(e -> e.getIndexOperator() != Expression.IndexOperator.ANN).collect(Collectors.toList());
QueryViewBuilder.QueryView queryView = new QueryViewBuilder(expressions, mergeRange).build();
Runnable onClose = () -> queryView.referencedIndexes.forEach(SSTableIndex::releaseQuietly);
KeyRangeIterator.Builder builder = command.rowFilter().isStrict()
? KeyRangeIntersectionIterator.builder(expressions.size(), onClose)
: KeyRangeUnionIterator.builder(expressions.size(), onClose);
? KeyRangeIntersectionIterator.builder(expressions.size(), queryView::close)
: KeyRangeUnionIterator.builder(expressions.size(), queryView::close);
try
{
@ -269,7 +320,7 @@ public class QueryController
return builder;
}
private void maybeTriggerGuardrails(QueryViewBuilder.QueryView queryView)
void maybeTriggerGuardrails(QueryViewBuilder.QueryView queryView)
{
int referencedIndexes = 0;
@ -313,117 +364,96 @@ public class QueryController
}
// This is an ANN only query
public KeyRangeIterator getTopKRows(RowFilter.Expression expression)
public CloseableIterator<PrimaryKeyWithScore> getTopKRows(QueryViewBuilder.QueryExpressionView queryExpressionView)
{
assert expression.operator() == Operator.ANN;
StorageAttachedIndex index = indexFor(expression);
assert index != null;
Expression planExpression = Expression.create(index).add(Operator.ANN, expression.getIndexValue().duplicate());
QueryViewBuilder.QueryView queryView = new QueryViewBuilder(Collections.singleton(planExpression), mergeRange).build();
Runnable onClose = () -> queryView.referencedIndexes.forEach(SSTableIndex::releaseQuietly);
assert queryExpressionView.expression.operator == Expression.IndexOperator.ANN;
List<CloseableIterator<PrimaryKeyWithScore>> intermediateResults = new ArrayList<>();
try
{
List<KeyRangeIterator> memtableResults = queryView.view
.stream()
.map(v -> v.memtableIndexes)
.flatMap(Collection::stream)
.map(idx -> idx.search(queryContext, planExpression, mergeRange))
.collect(Collectors.toList());
List<KeyRangeIterator> sstableIntersections = queryView.view
.stream()
.map(this::createRowIdIterator)
.collect(Collectors.toList());
return IndexSearchResultIterator.build(sstableIntersections, memtableResults, queryView.referencedIndexes, queryContext, onClose);
for (MemtableIndex memtableIndex : queryExpressionView.memtableIndexes)
intermediateResults.add(memtableIndex.orderBy(queryContext, queryExpressionView.expression, mergeRange));
for (SSTableIndex sstableIndex : queryExpressionView.sstableIndexes)
intermediateResults.addAll(sstableIndex.orderBy(queryExpressionView.expression, mergeRange, queryContext));
return intermediateResults.isEmpty() ? CloseableIterator.empty()
: new MergePrimaryKeyWithScoreIterator(intermediateResults);
}
catch (Throwable t)
{
// all sstable indexes in view have been referenced, need to clean up when exception is thrown
onClose.run();
throw t;
queryExpressionView.sstableIndexes.forEach(SSTableIndex::releaseQuietly);
intermediateResults.forEach(FileUtils::closeQuietly);
throw Throwables.cleaned(t);
}
}
// This is a hybrid query. We apply all other predicates before ordering and limiting.
public KeyRangeIterator getTopKRows(KeyRangeIterator source, RowFilter.Expression expression)
public CloseableIterator<PrimaryKeyWithScore> getTopKRows(KeyRangeIterator source, QueryViewBuilder.QueryExpressionView queryExpressionView)
{
return new KeyRangeOrderingIterator(source, orderChunkSize, list -> this.getTopKRows(list, expression));
List<PrimaryKey> primaryKeys = materializeKeysAndCloseSource(source);
if (primaryKeys == null)
return getTopKRows(queryExpressionView);
if (primaryKeys.isEmpty())
return CloseableIterator.empty();
return getTopKRows(primaryKeys, queryExpressionView);
}
private KeyRangeIterator getTopKRows(List<PrimaryKey> rawSourceKeys, RowFilter.Expression expression)
private CloseableIterator<PrimaryKeyWithScore> getTopKRows(List<PrimaryKey> sourceKeys, QueryViewBuilder.QueryExpressionView queryExpressionView)
{
VectorQueryContext vectorQueryContext = queryContext.vectorContext();
// Filter out PKs now. Each PK is passed to every segment of the ANN index, so filtering shadowed keys
// eagerly can save some work when going from PK to row id for on disk segments.
// Since the result is shared with multiple streams, we use an unmodifiable list.
List<PrimaryKey> sourceKeys = rawSourceKeys.stream().filter(vectorQueryContext::shouldInclude).collect(Collectors.toList());
StorageAttachedIndex index = indexFor(expression);
assert index != null : "Cannot do ANN ordering on an unindexed column";
Expression planExpression = Expression.create(index);
planExpression.add(Operator.ANN, expression.getIndexValue().duplicate());
QueryViewBuilder.QueryView queryView = new QueryViewBuilder(Collections.singleton(planExpression), mergeRange).build();
Runnable onClose = () -> queryView.referencedIndexes.forEach(SSTableIndex::releaseQuietly);
List<CloseableIterator<PrimaryKeyWithScore>> intermediateResults = new ArrayList<>();
try
{
List<KeyRangeIterator> memtableResults = queryView.view
.stream()
.map(v -> v.memtableIndexes)
.flatMap(Collection::stream)
.map(idx -> idx.limitToTopResults(sourceKeys, planExpression, vectorQueryContext.limit()))
.collect(Collectors.toList());
List<KeyRangeIterator> sstableIntersections = queryView.view
.stream()
.flatMap(pair -> pair.sstableIndexes.stream())
.map(idx -> {
try
{
return idx.limitToTopKResults(queryContext, sourceKeys, planExpression);
}
catch (IOException e)
{
throw new UncheckedIOException(e);
}
})
.collect(Collectors.toList());
return IndexSearchResultIterator.build(sstableIntersections, memtableResults, queryView.referencedIndexes, queryContext, onClose);
for (MemtableIndex memtableIndex : queryExpressionView.memtableIndexes)
intermediateResults.add(memtableIndex.orderResultsBy(queryContext, sourceKeys, queryExpressionView.expression));
for (SSTableIndex sstableIndex : queryExpressionView.sstableIndexes)
intermediateResults.addAll(sstableIndex.orderResultsBy(queryContext, sourceKeys, queryExpressionView.expression));
return intermediateResults.isEmpty() ? CloseableIterator.empty()
: new MergePrimaryKeyWithScoreIterator(intermediateResults);
}
catch (Throwable t)
{
// all sstable indexes in view have been referenced, need to clean up when exception is thrown
onClose.run();
throw t;
queryExpressionView.sstableIndexes.forEach(SSTableIndex::releaseQuietly);
intermediateResults.forEach(FileUtils::closeQuietly);
throw Throwables.cleaned(t);
}
}
/**
* Create row id iterator from different indexes' on-disk searcher of the same sstable
* Materialize the keys from the given source iterator. If there is a meaningful {@link #mergeRange}, the keys
* are filtered to only include those within the range. Note: closes the source iterator.
* @param source The source iterator to fully consume by materializing its keys
* @return The list of materialized keys within the {@link #mergeRange}, or return null if source exceeded the
* materialized keys limit.
*/
private KeyRangeIterator createRowIdIterator(QueryViewBuilder.QueryExpressionView indexExpression)
private List<PrimaryKey> materializeKeysAndCloseSource(KeyRangeIterator source)
{
List<KeyRangeIterator> subIterators = indexExpression.sstableIndexes
.stream()
.map(index ->
{
try
{
List<KeyRangeIterator> iterators = index.search(indexExpression.expression, mergeRange, queryContext);
// concat the result from multiple segments for the same index
return KeyRangeConcatIterator.builder(iterators.size()).add(iterators).build();
}
catch (Throwable ex)
{
throw Throwables.cleaned(ex);
}
}).collect(Collectors.toList());
try (source)
{
// Skip to the first key (which is really just a token) in the range if it is not the minimum token
if (!mergeRange.left.isMinimum())
source.skipTo(firstPrimaryKey);
return KeyRangeUnionIterator.build(subIterators);
if (!source.hasNext())
return List.of();
PrimaryKey maxToken = keyFactory.create(mergeRange.right.getToken());
boolean hasLimitingMaxToken = !maxToken.token().isMinimum() && maxToken.compareTo(source.getMaximum()) < 0;
List<PrimaryKey> primaryKeys = new ArrayList<>();
int count = 0;
while (source.hasNext())
{
PrimaryKey next = source.next();
if (hasLimitingMaxToken && next.compareTo(maxToken) > 0)
break;
primaryKeys.add(next);
if (MAX_MATERIALIZED_KEYS < ++count)
{
Tracing.trace("WHERE clause generated more than {} rows. Switching to ORDER BY then post filter.", MAX_MATERIALIZED_KEYS);
return null;
}
}
return primaryKeys;
}
}
// Note: This method assumes that the selects method has already been called for the

29
src/java/org/apache/cassandra/index/sai/plan/QueryMaterializesTooManyPrimaryKeysException.java Normal file
View File

@ -0,0 +1,29 @@
/*
* 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 org.apache.cassandra.db.RejectException;
public class QueryMaterializesTooManyPrimaryKeysException extends RejectException
{
public QueryMaterializesTooManyPrimaryKeysException(String msg)
{
super(msg);
}
}

19
src/java/org/apache/cassandra/index/sai/plan/QueryViewBuilder.java
View File

@ -25,7 +25,9 @@ import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.memtable.Memtable;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.disk.SSTableIndex;
import org.apache.cassandra.index.sai.memory.MemtableIndex;
@ -64,9 +66,18 @@ public class QueryViewBuilder
this.memtableIndexes = memtableIndexes;
this.sstableIndexes = sstableIndexes;
}
public ColumnFamilyStore.ViewFragment computeViewFragment()
{
// Because the SSTableIndex holds a reference to the SSTableReader, we know the sstable is still accessible
// so it is safe to build a view fragment.
List<Memtable> memtables = memtableIndexes.stream().map(MemtableIndex::getMemtable).collect(Collectors.toList());
List<SSTableReader> sstableReaders = sstableIndexes.stream().map(SSTableIndex::getSSTable).collect(Collectors.toList());
return new ColumnFamilyStore.ViewFragment(sstableReaders, memtables);
}
}
public static class QueryView
public static class QueryView implements AutoCloseable
{
public final Collection<QueryExpressionView> view;
public final Set<SSTableIndex> referencedIndexes;
@ -76,6 +87,12 @@ public class QueryViewBuilder
this.view = view;
this.referencedIndexes = referencedIndexes;
}
@Override
public void close()
{
referencedIndexes.forEach(SSTableIndex::releaseQuietly);
}
}
protected QueryView build()

2
src/java/org/apache/cassandra/index/sai/plan/StorageAttachedIndexQueryPlan.java
View File

@ -150,7 +150,7 @@ public class StorageAttachedIndexQueryPlan implements Index.QueryPlan
return partitions -> partitions;
// in case of top-k query, filter out rows that are not actually global top-K
return partitions -> (PartitionIterator) new VectorTopKProcessor(command).filter(partitions);
return partitions -> (PartitionIterator) new VectorTopKProcessor(command).consumeSortByScoreAndTakeTopK(partitions);
}
/**

492
src/java/org/apache/cassandra/index/sai/plan/StorageAttachedIndexSearcher.java
View File

@ -19,19 +19,26 @@
package org.apache.cassandra.index.sai.plan;
import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.PriorityQueue;
import java.util.Queue;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import io.netty.util.concurrent.FastThreadLocal;
import org.apache.cassandra.cql3.Operator;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.ClusteringBound;
import org.apache.cassandra.db.ClusteringComparator;
@ -59,13 +66,18 @@ import org.apache.cassandra.dht.Token;
import org.apache.cassandra.exceptions.RequestTimeoutException;
import org.apache.cassandra.index.Index;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.metrics.TableQueryMetrics;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.RangeUtil;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.utils.AbstractIterator;
import org.apache.cassandra.utils.Clock;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.FBUtilities;
public class StorageAttachedIndexSearcher implements Index.Searcher
{
@ -120,27 +132,64 @@ public class StorageAttachedIndexSearcher implements Index.Searcher
public UnfilteredPartitionIterator search(ReadExecutionController executionController) throws RequestTimeoutException
{
if (!command.isTopK())
return new ResultRetriever(executionController, false);
{
return new ResultRetriever(executionController);
}
else
{
Supplier<ResultRetriever> resultSupplier = () -> new ResultRetriever(executionController, true);
// VSTODO performance: if there is shadowed primary keys, we have to at least query twice.
// First time to find out there are shadow keys, second time to find out there are no more shadow keys.
while (true)
// Need a consistent view of the memtables/sstables and their associated index, so we get the view now
// and propagate it as needed.
try (QueryViewBuilder.QueryView queryView = buildAnnQueryView())
{
long lastShadowedKeysCount = queryContext.vectorContext().getShadowedPrimaryKeys().size();
ResultRetriever result = resultSupplier.get();
UnfilteredPartitionIterator topK = (UnfilteredPartitionIterator) new VectorTopKProcessor(command).filter(result);
long currentShadowedKeysCount = queryContext.vectorContext().getShadowedPrimaryKeys().size();
if (lastShadowedKeysCount == currentShadowedKeysCount)
return topK;
queryController.maybeTriggerGuardrails(queryView);
ScoreOrderedResultRetriever result = new ScoreOrderedResultRetriever(executionController, queryView);
// takeTopKThenSortByPrimaryKey eagerly consumes up to k rows from the result because search must
// produce an iterator in PrimaryKey order.
return (UnfilteredPartitionIterator) new VectorTopKProcessor(command).takeTopKThenSortByPrimaryKey(result);
}
}
}
private class ResultRetriever extends AbstractIterator<UnfilteredRowIterator> implements UnfilteredPartitionIterator
private QueryViewBuilder.QueryView buildAnnQueryView()
{
RowFilter.Expression annExpression = null;
for (RowFilter.Expression expression : queryController.indexFilter().getExpressions())
{
if (expression.operator() == Operator.ANN)
{
if (annExpression != null)
throw new IllegalStateException("Multiple ANN expressions in a single query are not supported");
annExpression = expression;
}
}
if (annExpression == null)
throw new IllegalStateException("No ANN expression found in query");
StorageAttachedIndex index = queryController.indexFor(annExpression);
Expression planExpression = Expression.create(index).add(Operator.ANN, annExpression.getIndexValue().duplicate());
return new QueryViewBuilder(Collections.singleton(planExpression), queryController.mergeRange()).build();
}
private abstract class AbstractRetreiver extends AbstractIterator<UnfilteredRowIterator> implements UnfilteredPartitionIterator
{
final FilterTree filterTree;
final ReadExecutionController executionController;
AbstractRetreiver(ReadExecutionController executionController)
{
this.executionController = executionController;
this.filterTree = Operation.buildFilter(queryController, queryController.usesStrictFiltering());
}
@Override
public TableMetadata metadata()
{
return queryController.metadata();
}
}
private class ResultRetriever extends AbstractRetreiver
{
private final PrimaryKey firstPrimaryKey;
private final PrimaryKey lastPrimaryKey;
@ -149,26 +198,21 @@ public class StorageAttachedIndexSearcher implements Index.Searcher
private AbstractBounds<PartitionPosition> currentKeyRange;
private final KeyRangeIterator resultKeyIterator;
private final FilterTree filterTree;
private final ReadExecutionController executionController;
private final PrimaryKey.Factory keyFactory;
private final boolean topK;
private final int partitionRowBatchSize;
private PrimaryKey lastKey;
private ResultRetriever(ReadExecutionController executionController, boolean topK)
private ResultRetriever(ReadExecutionController executionController)
{
super(executionController);
this.keyRanges = queryController.dataRanges().iterator();
this.firstDataRange = keyRanges.next();
this.currentKeyRange = firstDataRange.keyRange();
this.resultKeyIterator = Operation.buildIterator(queryController);
this.filterTree = Operation.buildFilter(queryController, queryController.usesStrictFiltering());
this.executionController = executionController;
this.keyFactory = queryController.primaryKeyFactory();
this.firstPrimaryKey = queryController.firstPrimaryKeyInRange();
this.lastPrimaryKey = queryController.lastPrimaryKeyInRange();
this.topK = topK;
// Ensure we don't fetch larger batches than the provided LIMIT to avoid fetching keys we won't use:
this.partitionRowBatchSize = Math.min(PARTITION_ROW_BATCH_SIZE, command.limits().count());
@ -458,114 +502,18 @@ public class StorageAttachedIndexSearcher implements Index.Searcher
queryContext.partitionsRead++;
queryContext.checkpoint();
UnfilteredRowIterator filtered = filterPartition(keys, partition, filterTree);
List<Row> filtered = filterPartition(partition, filterTree, queryContext);
// Note that we record the duration of the read after post-filtering, which actually
// materializes the rows from disk.
tableQueryMetrics.postFilteringReadLatency.update(Clock.Global.nanoTime() - startTimeNanos, TimeUnit.NANOSECONDS);
return filtered;
return filtered != null
? new SinglePartitionIterator(partition, partition.staticRow(), filtered.iterator())
: null;
}
}
private UnfilteredRowIterator filterPartition(List<PrimaryKey> keys, UnfilteredRowIterator partition, FilterTree tree)
{
Row staticRow = partition.staticRow();
DecoratedKey partitionKey = partition.partitionKey();
List<Unfiltered> matches = new ArrayList<>();
boolean hasMatch = false;
Set<PrimaryKey> keysToShadow = topK ? new HashSet<>(keys) : Collections.emptySet();
while (partition.hasNext())
{
Unfiltered unfiltered = partition.next();
if (unfiltered.isRow())
{
queryContext.rowsFiltered++;
if (tree.isSatisfiedBy(partitionKey, (Row) unfiltered, staticRow))
{
matches.add(unfiltered);
hasMatch = true;
if (topK)
{
PrimaryKey shadowed = keyFactory.hasClusteringColumns()
? keyFactory.create(partitionKey, ((Row) unfiltered).clustering())
: keyFactory.create(partitionKey);
keysToShadow.remove(shadowed);
}
}
}
}
// If any non-static rows match the filter, there should be no need to shadow the static primary key:
if (topK && hasMatch && keyFactory.hasClusteringColumns())
keysToShadow.remove(keyFactory.create(partitionKey, Clustering.STATIC_CLUSTERING));
// We may not have any non-static row data to filter...
if (!hasMatch)
{
queryContext.rowsFiltered++;
if (tree.isSatisfiedBy(partitionKey, staticRow, staticRow))
{
hasMatch = true;
if (topK)
keysToShadow.clear();
}
}
if (topK && !keysToShadow.isEmpty())
{
// Record primary keys shadowed by expired TTLs, row tombstones, or range tombstones:
queryContext.vectorContext().recordShadowedPrimaryKeys(keysToShadow);
}
if (!hasMatch)
{
// If there are no matches, return an empty partition. If reconciliation is required at the
// coordinator, replica filtering protection may make a second round trip to complete its view
// of the partition.
return null;
}
// Return all matches found, along with the static row...
return new PartitionIterator(partition, staticRow, matches.iterator());
}
private class PartitionIterator extends AbstractUnfilteredRowIterator
{
private final Iterator<Unfiltered> rows;
public PartitionIterator(UnfilteredRowIterator partition, Row staticRow, Iterator<Unfiltered> rows)
{
super(partition.metadata(),
partition.partitionKey(),
partition.partitionLevelDeletion(),
partition.columns(),
staticRow,
partition.isReverseOrder(),
partition.stats());
this.rows = rows;
}
@Override
protected Unfiltered computeNext()
{
return rows.hasNext() ? rows.next() : endOfData();
}
}
@Override
public TableMetadata metadata()
{
return queryController.metadata();
}
@Override
public void close()
{
@ -574,6 +522,306 @@ public class StorageAttachedIndexSearcher implements Index.Searcher
}
}
private static List<Row> filterPartition(UnfilteredRowIterator partition, FilterTree tree, QueryContext context)
{
Row staticRow = partition.staticRow();
DecoratedKey partitionKey = partition.partitionKey();
List<Row> matches = new ArrayList<>();
boolean hasMatch = false;
while (partition.hasNext())
{
Unfiltered unfiltered = partition.next();
if (unfiltered.isRow())
{
context.rowsFiltered++;
if (tree.isSatisfiedBy(partitionKey, (Row) unfiltered, staticRow))
{
matches.add((Row) unfiltered);
hasMatch = true;
}
}
}
// We may not have any non-static row data to filter...
if (!hasMatch)
{
context.rowsFiltered++;
if (tree.isSatisfiedBy(partitionKey, staticRow, staticRow))
{
hasMatch = true;
}
}
if (!hasMatch)
{
// If there are no matches, return an empty partition. If reconciliation is required at the
// coordinator, replica filtering protection may make a second round trip to complete its view
// of the partition.
return null;
}
// Return all matches found
return matches;
}
private static class SinglePartitionIterator extends AbstractUnfilteredRowIterator
{
private final Iterator<Row> rows;
public SinglePartitionIterator(UnfilteredRowIterator partition, Row staticRow, Iterator<Row> rows)
{
super(partition.metadata(),
partition.partitionKey(),
partition.partitionLevelDeletion(),
partition.columns(),
staticRow,
partition.isReverseOrder(),
partition.stats());
this.rows = rows;
}
@Override
protected Unfiltered computeNext()
{
return rows.hasNext() ? rows.next() : endOfData();
}
}
/**
* A result retriever that consumes an iterator primary keys sorted by some score, materializes the row for each
* primary key (currently, each primary key is required to be fully qualified and should only point to one row),
* apply the filter tree to the row to test that the real row satisfies the WHERE clause, and finally tests
* that the row is valid for the ORDER BY clause. The class performs some optimizations to avoid materializing
* rows unnecessarily. See the class for more details.
* <p>
* The resulting {@link UnfilteredRowIterator} objects are not guaranteed to be in any particular order. It is
* the responsibility of the caller to sort the results if necessary.
*/
public class ScoreOrderedResultRetriever extends AbstractRetreiver
{
private final ColumnFamilyStore.ViewFragment view;
private final List<AbstractBounds<PartitionPosition>> keyRanges;
private final boolean coversFullRing;
private final CloseableIterator<PrimaryKeyWithScore> scoredPrimaryKeyIterator;
private final boolean isVectorColumnStatic;
private final HashSet<PrimaryKey> processedKeys;
private final Queue<UnfilteredRowIterator> pendingRows;
// The limit requested by the query. We cannot load more than softLimit rows in bulk because we only want
// to fetch the topk rows where k is the limit. However, we allow the iterator to fetch more rows than the
// soft limit to avoid confusing behavior. When the softLimit is reached, the iterator will fetch one row
// at a time.
private final int softLimit;
private int returnedRowCount = 0;
private ScoreOrderedResultRetriever(ReadExecutionController executionController,
QueryViewBuilder.QueryView queryView)
{
super(executionController);
assert queryView.view.size() == 1;
QueryViewBuilder.QueryExpressionView queryExpressionView = queryView.view.stream().findFirst().get();
this.view = queryExpressionView.computeViewFragment();
this.keyRanges = queryController.dataRanges().stream().map(DataRange::keyRange).collect(Collectors.toList());
this.coversFullRing = keyRanges.size() == 1 && RangeUtil.coversFullRing(keyRanges.get(0));
this.scoredPrimaryKeyIterator = Operation.buildIteratorForOrder(queryController, queryExpressionView);
this.isVectorColumnStatic = queryExpressionView.expression.getIndexTermType().columnMetadata().isStatic();
this.softLimit = command.limits().count();
this.processedKeys = new HashSet<>(softLimit);
this.pendingRows = new ArrayDeque<>(softLimit);
}
@Override
public UnfilteredRowIterator computeNext()
{
if (pendingRows.isEmpty())
fillPendingRows();
returnedRowCount++;
// Because we know ordered keys are fully qualified, we do not iterate partitions
return !pendingRows.isEmpty() ? pendingRows.poll() : endOfData();
}
/**
* Fills the pendingRows queue to generate a queue of row iterators for the supplied keys by repeatedly calling
* {@link #readAndValidatePartition} until it gives enough non-null results.
*/
private void fillPendingRows()
{
// Group PKs by source sstable/memtable
Map<PrimaryKey, List<PrimaryKeyWithScore>> groupedKeys = new HashMap<>();
// We always want to get at least 1. When the vector column is static, we cannot batch because we need to
// retain the score ordering a bit longer.
int rowsToRetrieve = isVectorColumnStatic ? 1 : Math.max(1, softLimit - returnedRowCount);
// We want to get the first unique `rowsToRetrieve` keys to materialize
// Don't pass the priority queue here because it is more efficient to add keys in bulk
fillKeys(groupedKeys, rowsToRetrieve, null);
// Sort the primary keys by PrK order, just in case that helps with cache and disk efficiency
PriorityQueue<PrimaryKey> primaryKeyPriorityQueue = new PriorityQueue<>(groupedKeys.keySet());
// drain groupedKeys into pendingRows
while (!groupedKeys.isEmpty())
{
PrimaryKey pk = primaryKeyPriorityQueue.poll();
List<PrimaryKeyWithScore> sourceKeys = groupedKeys.remove(pk);
UnfilteredRowIterator partitionIterator = readAndValidatePartition(pk, sourceKeys);
if (partitionIterator != null)
pendingRows.add(partitionIterator);
else
// The current primaryKey did not produce a partition iterator. We know the caller will need
// `rowsToRetrieve` rows, so we get the next unique key and add it to the queue.
fillKeys(groupedKeys, 1, primaryKeyPriorityQueue);
}
}
/**
* Fills the `groupedKeys` Map with the next `count` unique primary keys that are in the keys produced by calling
* {@link #nextSelectedKeyInRange()}. We map PrimaryKey to a list of PrimaryKeyWithScore because the same
* primary key can be in the result set multiple times, but with different source tables.
* @param groupedKeys the map to fill
* @param count the number of unique PrimaryKeys to consume from the iterator
* @param primaryKeyPriorityQueue the priority queue to add new keys to. If the queue is null, we do not add
* keys to the queue.
*/
private void fillKeys(Map<PrimaryKey, List<PrimaryKeyWithScore>> groupedKeys, int count, PriorityQueue<PrimaryKey> primaryKeyPriorityQueue)
{
int initialSize = groupedKeys.size();
while (groupedKeys.size() - initialSize < count)
{
PrimaryKeyWithScore primaryKeyWithScore = nextSelectedKeyInRange();
if (primaryKeyWithScore == null)
return;
PrimaryKey nextPrimaryKey = primaryKeyWithScore.primaryKey();
List<PrimaryKeyWithScore> accumulator = groupedKeys.computeIfAbsent(nextPrimaryKey, k -> new ArrayList<>());
if (primaryKeyPriorityQueue != null && accumulator.isEmpty())
primaryKeyPriorityQueue.add(nextPrimaryKey);
accumulator.add(primaryKeyWithScore);
}
}
/**
* Determine if the key is in one of the queried key ranges. We do not iterate through results in
* {@link PrimaryKey} order, so we have to check each range.
* @param key the key to test
* @return true if the key is in one of the queried key ranges
*/
private boolean isInRange(DecoratedKey key)
{
if (coversFullRing)
return true;
for (AbstractBounds<PartitionPosition> range : keyRanges)
if (range.contains(key))
return true;
return false;
}
/**
* Returns the next available key contained by one of the keyRanges and selected by the queryController.
* If the next key falls out of the current key range, it skips to the next key range, and so on.
* If no more keys acceptd by the controller are available, returns null.
*/
private @Nullable PrimaryKeyWithScore nextSelectedKeyInRange()
{
while (scoredPrimaryKeyIterator.hasNext())
{
PrimaryKeyWithScore key = scoredPrimaryKeyIterator.next();
if (isInRange(key.primaryKey().partitionKey()) && !queryController.doesNotSelect(key.primaryKey()))
return key;
}
return null;
}
/**
* Reads and validates a partition for a given primary key against its sources.
* <p>
* @param pk The primary key of the partition to read and validate
* @param sourceKeys A list of PrimaryKeyWithScore objects associated with the primary key.
* Multiple sort keys can exist for the same primary key when data comes from different
* sstables or memtables.
*
* @return An UnfilteredRowIterator containing the validated partition data, or null if:
* - The key has already been processed
* - The partition does not pass index filters
* - The partition contains no valid rows
* - The row data does not match the index metadata for any of the provided primary keys
*/
public UnfilteredRowIterator readAndValidatePartition(PrimaryKey pk, List<PrimaryKeyWithScore> sourceKeys)
{
// If we've already processed the key, we can skip it. Because the score ordered iterator does not
// deduplicate rows, we could see dupes if a row is in the ordering index multiple times. This happens
// in the case of dupes and of overwrites.
if (processedKeys.contains(pk))
return null;
try (UnfilteredRowIterator partition = queryController.queryStorage(pk, view, executionController))
{
queryContext.partitionsRead++;
queryContext.checkpoint();
List<Row> clusters = filterPartition(partition, filterTree, queryContext);
if (clusters == null)
{
// Key counts as processed because the materialized row didn't satisfy the filter logic
processedKeys.add(pk);
return null;
}
Row staticRow = partition.staticRow();
long now = FBUtilities.nowInSeconds();
// If the pk is static, then we must check that the static row satisfies the source key's validity check.
// Otherwise, we need to make sure that we have one row in the cluster result and then we use that
// for checking validity.
Row representativeRow;
if (pk.kind() == PrimaryKey.Kind.STATIC)
{
representativeRow = staticRow;
}
else
{
if (clusters.isEmpty())
{
// Key counts as processed because the materialized row didn't satisfy the filter logic
processedKeys.add(pk);
return null;
}
representativeRow = clusters.get(0);
assert clusters.size() == 1 : "Expect 1 result row, but got: " + clusters.size();
}
// Each of sourceKeys are equal with respect to primary key equality, but they have different source tables.
// As long as one is valid, we consider the row valid.
for (PrimaryKeyWithScore sourceKey : sourceKeys)
{
assert sourceKey.primaryKey().kind() == pk.kind();
if (sourceKey.isIndexDataValid(representativeRow, now))
{
processedKeys.add(pk);
return new SinglePartitionIterator(partition, staticRow, clusters.iterator());
}
}
// Key does not count as processed because the only thing that "failed" is the validity check on the
// grouped source keys, and it is possible that the score ordered iterator has the same key in the
// iterator lower. We only get here when a vector's value is updated to a more distant vector, so
// the old value ranks high in the iterator, but isn't the current value for the materialized row.
return null;
}
}
public void close()
{
FileUtils.closeQuietly(scoredPrimaryKeyIterator);
}
}
/**
* Used by {@link StorageAttachedIndexSearcher#filterReplicaFilteringProtection} to filter rows for columns that
* have transformations so won't get handled correctly by the row filter.

49
src/java/org/apache/cassandra/index/sai/plan/VectorTopKProcessor.java
View File

@ -85,7 +85,7 @@ public class VectorTopKProcessor
* Filter given partitions and keep the rows with the highest scores. In case of {@link UnfilteredPartitionIterator},
* all tombstones will be kept.
*/
public <U extends Unfiltered, R extends BaseRowIterator<U>, P extends BasePartitionIterator<R>> BasePartitionIterator<?> filter(P partitions)
public <U extends Unfiltered, R extends BaseRowIterator<U>, P extends BasePartitionIterator<R>> BasePartitionIterator<?> consumeSortByScoreAndTakeTopK(P partitions)
{
// priority queue ordered by score in ascending order
PriorityQueue<Triple<PartitionInfo, Row, Float>> topK = new PriorityQueue<>(limit + 1, Comparator.comparing(Triple::getRight));
@ -161,6 +161,53 @@ public class VectorTopKProcessor
return 0;
}
/**
* Filter given partitions and keep the rows with the highest scores. In case of {@link UnfilteredPartitionIterator},
* all tombstones will be kept.
*/
public <U extends Unfiltered, R extends BaseRowIterator<U>, P extends BasePartitionIterator<R>> BasePartitionIterator<?> takeTopKThenSortByPrimaryKey(P partitions)
{
try (partitions)
{
TreeMap<PartitionInfo, TreeSet<Unfiltered>> unfilteredByPartition = new TreeMap<>(Comparator.comparing(pi -> pi.key));
int rowsMatched = 0;
while (rowsMatched < limit && partitions.hasNext())
{
try (BaseRowIterator<?> partitionRowIterator = partitions.next())
{
rowsMatched += processSingleRowPartition(unfilteredByPartition, partitionRowIterator, limit - rowsMatched);
}
}
return new InMemoryUnfilteredPartitionIterator(command, unfilteredByPartition);
}
}
/**
* Processes a single partition, without scoring it.
*/
private int processSingleRowPartition(TreeMap<PartitionInfo, TreeSet<Unfiltered>> unfilteredByPartition,
BaseRowIterator<?> partitionRowIterator,
int reamining)
{
if (!partitionRowIterator.hasNext())
return 0;
// Always include tombstones for coordinator. It relies on ReadCommand#withMetricsRecording to throw
// TombstoneOverwhelmingException to prevent OOM.
PartitionInfo partitionInfo = PartitionInfo.create(partitionRowIterator);
TreeSet<Unfiltered> map = unfilteredByPartition.computeIfAbsent(partitionInfo, k -> new TreeSet<>(command.metadata().comparator));
int added = 0;
while (partitionRowIterator.hasNext() && added < reamining)
{
Unfiltered unfiltered = partitionRowIterator.next();
map.add(unfiltered);
if (unfiltered.isRow())
added++;
}
return added;
}
private Pair<StorageAttachedIndex, float[]> findTopKIndex()
{

246
src/java/org/apache/cassandra/index/sai/utils/CellWithSource.java Normal file
View File

@ -0,0 +1,246 @@
/*
* 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.nio.ByteBuffer;
import org.apache.cassandra.db.CellSourceIdentifier;
import org.apache.cassandra.db.DeletionPurger;
import org.apache.cassandra.db.Digest;
import org.apache.cassandra.db.marshal.ValueAccessor;
import org.apache.cassandra.db.memtable.Memtable;
import org.apache.cassandra.db.rows.Cell;
import org.apache.cassandra.db.rows.CellPath;
import org.apache.cassandra.db.rows.ColumnData;
import org.apache.cassandra.db.rows.ComplexColumnData;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.utils.ObjectSizes;
import org.apache.cassandra.utils.memory.ByteBufferCloner;
/**
* A wrapped {@link Cell} that includes a reference to the cell's source table via {@link CellSourceIdentifier}
* @param <T> the type of the cell's value
*/
public class CellWithSource<T> extends Cell<T>
{
private static final long EMPTY_SIZE = ObjectSizes.measure(new CellWithSource<>(null, null, null));
private final Cell<T> cell;
private final CellSourceIdentifier source;
public CellWithSource(Cell<T> cell, CellSourceIdentifier source)
{
this(cell.column(), cell, source);
assert source instanceof Memtable || source instanceof SSTableId : "Source has unexpected type: " + (source == null ? "null" : source.getClass());
}
private CellWithSource(ColumnMetadata column, Cell<T> cell, CellSourceIdentifier source)
{
super(column);
this.cell = cell;
this.source = source;
}
public CellSourceIdentifier sourceTable()
{
return source;
}
@Override
public boolean isCounterCell()
{
return cell.isCounterCell();
}
@Override
public T value()
{
return cell.value();
}
@Override
public ValueAccessor<T> accessor()
{
return cell.accessor();
}
@Override
public long timestamp()
{
return cell.timestamp();
}
@Override
public int ttl()
{
return cell.ttl();
}
@Override
public long localDeletionTime()
{
return cell.localDeletionTime();
}
@Override
public boolean isTombstone()
{
return cell.isTombstone();
}
@Override
public boolean isExpiring()
{
return cell.isExpiring();
}
@Override
public boolean isLive(long nowInSec)
{
return cell.isLive(nowInSec);
}
@Override
public CellPath path()
{
return cell.path();
}
@Override
public Cell<?> withUpdatedColumn(ColumnMetadata newColumn)
{
return wrapIfNew(cell.withUpdatedColumn(newColumn));
}
@Override
public Cell<?> withUpdatedValue(ByteBuffer newValue)
{
return wrapIfNew(cell.withUpdatedValue(newValue));
}
@Override
public Cell<?> withUpdatedTimestampAndLocalDeletionTime(long newTimestamp, long newLocalDeletionTime)
{
return wrapIfNew(cell.withUpdatedTimestampAndLocalDeletionTime(newTimestamp, newLocalDeletionTime));
}
@Override
public Cell<?> withSkippedValue()
{
return wrapIfNew(cell.withSkippedValue());
}
@Override
public Cell<?> clone(ByteBufferCloner cloner)
{
return wrapIfNew(cell.clone(cloner));
}
@Override
public int dataSize()
{
return cell.dataSize();
}
@Override
public long unsharedHeapSizeExcludingData()
{
return cell.unsharedHeapSizeExcludingData();
}
@Override
public long unsharedHeapSize()
{
return cell.unsharedHeapSize() + EMPTY_SIZE;
}
@Override
public void validate()
{
cell.validate();
}
@Override
public boolean hasInvalidDeletions()
{
return cell.hasInvalidDeletions();
}
@Override
public void digest(Digest digest)
{
cell.digest(digest);
}
@Override
public ColumnData updateAllTimestamp(long newTimestamp)
{
ColumnData maybeNewCell = cell.updateAllTimestamp(newTimestamp);
if (maybeNewCell instanceof Cell)
return wrapIfNew((Cell<?>) maybeNewCell);
if (maybeNewCell instanceof ComplexColumnData)
return ((ComplexColumnData) maybeNewCell).transform(this::wrapIfNew);
// It's not clear when we would hit this code path, but it seems we should not
// hit this from SAI.
throw new IllegalStateException("Expected a Cell instance, but got " + maybeNewCell);
}
@Override
public Cell<?> markCounterLocalToBeCleared()
{
return wrapIfNew(cell.markCounterLocalToBeCleared());
}
@Override
public Cell<?> purge(DeletionPurger purger, long nowInSec)
{
return wrapIfNew(cell.purge(purger, nowInSec));
}
@Override
public Cell<?> purgeDataOlderThan(long timestamp)
{
return wrapIfNew(cell.purgeDataOlderThan(timestamp));
}
@Override
protected int localDeletionTimeAsUnsignedInt()
{
// Cannot call cell's localDeletionTimeAsUnsignedInt() because it's protected.
throw new UnsupportedOperationException();
}
@Override
public long maxTimestamp()
{
return cell.maxTimestamp();
}
private Cell<?> wrapIfNew(Cell<?> maybeNewCell)
{
if (maybeNewCell == null)
return null;
// If the cell's method returned a reference to the same cell, then
// we can skip creating a new wrapper.
if (maybeNewCell == this.cell)
return this;
return new CellWithSource<>(maybeNewCell, source);
}
}

70
src/java/org/apache/cassandra/index/sai/utils/MergePrimaryKeyWithScoreIterator.java Normal file
View File

@ -0,0 +1,70 @@
/*
* 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.Collection;
import java.util.PriorityQueue;
import com.google.common.collect.Iterators;
import com.google.common.collect.PeekingIterator;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.utils.AbstractIterator;
import org.apache.cassandra.utils.CloseableIterator;
// TODO: this implementation is sub-optimal due to the combination of PriorityQueue poll/add. A non reducing version of
// the MergeIterator would be better
public class MergePrimaryKeyWithScoreIterator extends AbstractIterator<PrimaryKeyWithScore>
{
private final PriorityQueue<PeekingIterator<PrimaryKeyWithScore>> queue;
private final Collection<CloseableIterator<PrimaryKeyWithScore>> iteratorsToClose;
public MergePrimaryKeyWithScoreIterator(Collection<CloseableIterator<PrimaryKeyWithScore>> iterators)
{
assert !iterators.isEmpty();
iteratorsToClose = iterators;
queue = new PriorityQueue<>(iterators.size(), (a, b) -> a.peek().compareTo(b.peek()));
for (CloseableIterator<PrimaryKeyWithScore> iterator : iterators)
{
if (iterator.hasNext())
queue.add(Iterators.peekingIterator(iterator));
}
}
@Override
protected PrimaryKeyWithScore computeNext()
{
if (queue.isEmpty())
return endOfData();
PeekingIterator<PrimaryKeyWithScore> iterator = queue.poll();
PrimaryKeyWithScore next = iterator.next();
if (iterator.hasNext())
queue.add(iterator);
return next;
}
@Override
public void close()
{
for (CloseableIterator<PrimaryKeyWithScore> iterator : iteratorsToClose)
FileUtils.closeQuietly(iterator);
}
}

392
src/java/org/apache/cassandra/index/sai/utils/RowWithSource.java Normal file
View File

@ -0,0 +1,392 @@
/*
* 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.Collection;
import java.util.Comparator;
import java.util.Iterator;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
import com.google.common.collect.Collections2;
import com.google.common.collect.Iterables;
import com.google.common.collect.Iterators;
import org.apache.cassandra.db.CellSourceIdentifier;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.DeletionPurger;
import org.apache.cassandra.db.DeletionTime;
import org.apache.cassandra.db.Digest;
import org.apache.cassandra.db.LivenessInfo;
import org.apache.cassandra.db.filter.ColumnFilter;
import org.apache.cassandra.db.memtable.Memtable;
import org.apache.cassandra.db.rows.Cell;
import org.apache.cassandra.db.rows.CellPath;
import org.apache.cassandra.db.rows.ColumnData;
import org.apache.cassandra.db.rows.ComplexColumnData;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.utils.BiLongAccumulator;
import org.apache.cassandra.utils.LongAccumulator;
import org.apache.cassandra.utils.ObjectSizes;
import org.apache.cassandra.utils.SearchIterator;
import org.apache.cassandra.utils.memory.Cloner;
/**
* A Row wrapper that has a {@link CellSourceIdentifier} that gets added to cell as part of the
* {@link #getCell(ColumnMetadata)} and {@link #getCell(ColumnMetadata, CellPath)} calls. This class
* can only be initiallized validly when all the cells share a common {@link CellSourceIdentifier}.
*/
public class RowWithSource implements Row
{
private static final long EMPTY_SIZE = ObjectSizes.measure(new RowWithSource(null, null));
private final Row row;
private final CellSourceIdentifier source;
public RowWithSource(Row row, CellSourceIdentifier source)
{
assert source instanceof Memtable || source instanceof SSTableId || (source == null && row == null) : "Expected Memtable or SSTableId, got " + source;
this.row = row;
this.source = source;
}
@Override
public Kind kind()
{
return row.kind();
}
@Override
public Clustering<?> clustering()
{
return row.clustering();
}
@Override
public void digest(Digest digest)
{
row.digest(digest);
}
@Override
public void validateData(TableMetadata metadata)
{
row.validateData(metadata);
}
@Override
public boolean hasInvalidDeletions()
{
return row.hasInvalidDeletions();
}
@Override
public Collection<ColumnMetadata> columns()
{
return row.columns();
}
@Override
public int columnCount()
{
return row.columnCount();
}
@Override
public Deletion deletion()
{
return row.deletion();
}
@Override
public LivenessInfo primaryKeyLivenessInfo()
{
return row.primaryKeyLivenessInfo();
}
@Override
public boolean isStatic()
{
return row.isStatic();
}
@Override
public boolean isEmpty()
{
return row.isEmpty();
}
@Override
public String toString(TableMetadata metadata)
{
return row.toString(metadata);
}
@Override
public boolean hasLiveData(long nowInSec, boolean enforceStrictLiveness)
{
return row.hasLiveData(nowInSec, enforceStrictLiveness);
}
@Override
public Cell<?> getCell(ColumnMetadata c)
{
Cell<?> cell = row.getCell(c);
if (cell == null)
return null;
return new CellWithSource<>(cell, source);
}
@Override
public Cell<?> getCell(ColumnMetadata c, CellPath path)
{
return wrapCell(row.getCell(c, path));
}
@Override
public ComplexColumnData getComplexColumnData(ColumnMetadata c)
{
return (ComplexColumnData) wrapColumnData(row.getComplexColumnData(c));
}
@Override
public ColumnData getColumnData(ColumnMetadata c)
{
return wrapColumnData(row.getColumnData(c));
}
@Override
public Iterable<Cell<?>> cells()
{
return Iterables.transform(row.cells(), this::wrapCell);
}
@Override
public Collection<ColumnData> columnData()
{
return Collections2.transform(row.columnData(), this::wrapColumnData);
}
@Override
public Iterable<Cell<?>> cellsInLegacyOrder(TableMetadata metadata, boolean reversed)
{
return Iterables.transform(row.cellsInLegacyOrder(metadata, reversed), this::wrapCell);
}
@Override
public boolean hasComplexDeletion()
{
return row.hasComplexDeletion();
}
@Override
public boolean hasComplex()
{
return row.hasComplex();
}
@Override
public boolean hasDeletion(long nowInSec)
{
return row.hasDeletion(nowInSec);
}
@Override
public SearchIterator<ColumnMetadata, ColumnData> searchIterator()
{
SearchIterator<ColumnMetadata, ColumnData> iterator = row.searchIterator();
return key -> wrapColumnData(iterator.next(key));
}
@Override
public Row filter(ColumnFilter filter, TableMetadata metadata)
{
return maybeWrapRow(row.filter(filter, metadata));
}
@Override
public Row filter(ColumnFilter filter, DeletionTime activeDeletion, boolean setActiveDeletionToRow, TableMetadata metadata)
{
return maybeWrapRow(row.filter(filter, activeDeletion, setActiveDeletionToRow, metadata));
}
@Override
public Row transformAndFilter(LivenessInfo info, Deletion deletion, Function<ColumnData, ColumnData> function)
{
return maybeWrapRow(row.transformAndFilter(info, deletion, function));
}
@Override
public Row transformAndFilter(Function<ColumnData, ColumnData> function)
{
return maybeWrapRow(row.transformAndFilter(function));
}
@Override
public Row clone(Cloner cloner)
{
return maybeWrapRow(row.clone(cloner));
}
@Override
public Row purgeDataOlderThan(long timestamp, boolean enforceStrictLiveness)
{
return maybeWrapRow(row.purgeDataOlderThan(timestamp, enforceStrictLiveness));
}
@Override
public Row purge(DeletionPurger purger, long nowInSec, boolean enforceStrictLiveness)
{
return maybeWrapRow(row.purge(purger, nowInSec, enforceStrictLiveness));
}
@Override
public Row withOnlyQueriedData(ColumnFilter filter)
{
return maybeWrapRow(row.withOnlyQueriedData(filter));
}
@Override
public Row markCounterLocalToBeCleared()
{
return maybeWrapRow(row.markCounterLocalToBeCleared());
}
@Override
public Row updateAllTimestamp(long newTimestamp)
{
return maybeWrapRow(row.updateAllTimestamp(newTimestamp));
}
@Override
public Row withRowDeletion(DeletionTime deletion)
{
return maybeWrapRow(row.withRowDeletion(deletion));
}
@Override
public int dataSize()
{
return row.dataSize();
}
@Override
public long unsharedHeapSize()
{
return row.unsharedHeapSize() + EMPTY_SIZE;
}
@Override
public long unsharedHeapSizeExcludingData()
{
return row.unsharedHeapSizeExcludingData() + EMPTY_SIZE;
}
@Override
public String toString(TableMetadata metadata, boolean fullDetails)
{
return row.toString(metadata, fullDetails);
}
@Override
public String toString(TableMetadata metadata, boolean includeClusterKeys, boolean fullDetails)
{
return row.toString(metadata, includeClusterKeys, fullDetails);
}
@Override
public void apply(Consumer<ColumnData> function)
{
row.apply(function);
}
@Override
public <A> void apply(BiConsumer<A, ColumnData> function, A arg)
{
row.apply(function, arg);
}
@Override
public long accumulate(LongAccumulator<ColumnData> accumulator, long initialValue)
{
return row.accumulate(accumulator, initialValue);
}
@Override
public long accumulate(LongAccumulator<ColumnData> accumulator, Comparator<ColumnData> comparator, ColumnData from, long initialValue)
{
return row.accumulate(accumulator, comparator, from, initialValue);
}
@Override
public <A> long accumulate(BiLongAccumulator<A, ColumnData> accumulator, A arg, long initialValue)
{
return row.accumulate(accumulator, arg, initialValue);
}
@Override
public <A> long accumulate(BiLongAccumulator<A, ColumnData> accumulator, A arg, Comparator<ColumnData> comparator, ColumnData from, long initialValue)
{
return row.accumulate(accumulator, arg, comparator, from, initialValue);
}
@Override
public Iterator<ColumnData> iterator()
{
return Iterators.transform(row.iterator(), this::wrapColumnData);
}
private ColumnData wrapColumnData(ColumnData c)
{
if (c == null)
return null;
if (c instanceof Cell<?>)
return new CellWithSource<>((Cell<?>) c, source);
if (c instanceof ComplexColumnData)
return ((ComplexColumnData) c).transform(c1 -> new CellWithSource<>(c1, source));
throw new IllegalStateException("Unexpected ColumnData type: " + c.getClass().getName());
}
private Cell<?> wrapCell(Cell<?> c)
{
return c != null ? new CellWithSource<>(c, source) : null;
}
private Row maybeWrapRow(Row r)
{
if (r == null)
return null;
if (r == this.row)
return this;
return new RowWithSource(r, source);
}
@Override
public String toString()
{
return "RowWithSourceTable{" +
row +
", source=" + source +
'}';
}
}

21
src/java/org/apache/cassandra/index/sai/view/IndexViewManager.java
View File

@ -167,13 +167,6 @@ public class IndexViewManager
continue;
}
if (sstableContext.indexDescriptor.isIndexEmpty(index.termType(), index.identifier()))
{
logger.debug(index.identifier().logMessage("No on-disk index was built for SSTable {} because the SSTable " +
"had no indexable rows for the index."), sstableContext.descriptor());
continue;
}
try
{
if (validation != IndexValidation.NONE)
@ -186,7 +179,19 @@ public class IndexViewManager
}
SSTableIndex ssTableIndex = sstableContext.newSSTableIndex(index);
logger.debug(index.identifier().logMessage("Successfully created index for SSTable {}."), sstableContext.descriptor());
// We used to skip these empty indexes. However, that leads to logically incomplete views of the table,
// so we keep them in the view now. For example, vector indexes use the view to materialize rows, and
// without a complete view, an sstable with no indexable vectors might still have valid data or
// tombstones necessary to ensure proper row materialization.
if (ssTableIndex.getRowCount() == 0)
{
logger.debug(index.identifier().logMessage("No on-disk index was built for SSTable {} because the SSTable " +
"had no indexable rows for the index."), sstableContext.descriptor());
}
else
{
logger.debug(index.identifier().logMessage("Successfully created index for SSTable {}."), sstableContext.descriptor());
}
// Try to add new index to the set, if set already has such index, we'll simply release and move on.
// This covers situation when SSTable collection has the same SSTable multiple

4
src/java/org/apache/cassandra/index/sai/view/View.java
View File

@ -52,7 +52,9 @@ public class View implements Iterable<SSTableIndex>
for (SSTableIndex sstableIndex : indexes)
{
this.view.put(sstableIndex.getSSTable().descriptor, sstableIndex);
if (!indexTermType.isVector())
// Skip vector indexes, since they are scatter gather for all terms. Skip empty indexes since they
// cannot be inserted into the tree due to the lack of min and max terms.
if (!indexTermType.isVector() && sstableIndex.getRowCount() > 0)
rangeTermTreeBuilder.add(sstableIndex);
}

6
src/java/org/apache/cassandra/index/sai/virtual/SSTableIndexesSystemView.java
View File

@ -102,6 +102,12 @@ public class SSTableIndexesSystemView extends AbstractVirtualTable
for (SSTableIndex sstableIndex : index.view())
{
// Empty indexes are tracked internally for the sake of having complete views. However,
// these indexes have not historically been exposed in this virtual table, so we skip
// them for now.
if (sstableIndex.getRowCount() == 0)
continue;
SSTableReader sstable = sstableIndex.getSSTable();
Descriptor descriptor = sstable.descriptor;
AbstractBounds<Token> bounds = sstable.getBounds();

2
src/java/org/apache/cassandra/index/sasi/SASIIndex.java
View File

@ -358,7 +358,7 @@ public class SASIIndex implements Index, INotificationConsumer
}
else if (notification instanceof MemtableSwitchedNotification)
{
index.switchMemtable(((MemtableSwitchedNotification) notification).memtable);
index.switchMemtable(((MemtableSwitchedNotification) notification).previous);
}
else if (notification instanceof MemtableDiscardedNotification)
{

5
src/java/org/apache/cassandra/io/sstable/SSTable.java
View File

@ -203,6 +203,11 @@ public abstract class SSTable
return descriptor.ksname;
}
public SSTableId getId()
{
return descriptor.id;
}
public List<String> getAllFilePaths()
{
List<String> ret = new ArrayList<>(components.size());

3
src/java/org/apache/cassandra/io/sstable/SSTableId.java
View File

@ -22,6 +22,7 @@ import java.nio.ByteBuffer;
import java.util.function.Supplier;
import java.util.stream.Stream;
import org.apache.cassandra.db.CellSourceIdentifier;
import org.apache.cassandra.io.util.File;
/**
@ -37,7 +38,7 @@ import org.apache.cassandra.io.util.File;
* - must be case-insensitive because the sstables can be stored on case-insensitive file system
* <p>
*/
public interface SSTableId
public interface SSTableId extends CellSourceIdentifier
{
/**
* Creates a byte format of the identifier that can be parsed by

8
src/java/org/apache/cassandra/notifications/MemtableSwitchedNotification.java
View File

@ -21,10 +21,12 @@ import org.apache.cassandra.db.memtable.Memtable;
public class MemtableSwitchedNotification implements INotification
{
public final Memtable memtable;
public final Memtable previous;
public final Memtable next;
public MemtableSwitchedNotification(Memtable switched)
public MemtableSwitchedNotification(Memtable switched, Memtable next)
{
this.memtable = switched;
this.previous = switched;
this.next = next;
}
}

36
test/unit/org/apache/cassandra/cql3/CQLTester.java
View File

@ -45,6 +45,7 @@ import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
@ -2549,6 +2550,41 @@ public abstract class CQLTester
return Arrays.asList(values);
}
/** @return a normalized vector with the given dimension */
public Vector<Float> randomVectorBoxed(int dimension)
{
float[] floats = randomVector(dimension);
return vector(floats);
}
public float[] randomVector(int dimension)
{
// this can be called from concurrent threads so don't use getRandom()
ThreadLocalRandom R = ThreadLocalRandom.current();
float[] vector = new float[dimension];
for (int i = 0; i < dimension; i++)
{
vector[i] = R.nextFloat();
}
normalize(vector);
return vector;
}
/** Normalize the given vector in-place */
protected static void normalize(float[] v)
{
float sum = 0.0f;
for (int i = 0; i < v.length; i++)
{
sum += v[i] * v[i];
}
sum = (float) Math.sqrt(sum);
for (int i = 0; i < v.length; i++)
v[i] /= sum;
}
@SafeVarargs
protected final <T> Vector<T> vector(T... values)
{

6
test/unit/org/apache/cassandra/db/lifecycle/TrackerTest.java
View File

@ -319,7 +319,8 @@ public class TrackerTest
tracker = cfs.getTracker();
listener = new MockListener(false);
tracker.subscribe(listener);
prev1 = tracker.switchMemtable(false, cfs.createMemtable(new AtomicReference<>(CommitLog.instance.getCurrentPosition())));
Memtable next1 = cfs.createMemtable(new AtomicReference<>(CommitLog.instance.getCurrentPosition()));
prev1 = tracker.switchMemtable(false, next1);
tracker.markFlushing(prev1);
reader = MockSchema.sstable(0, 10, true, cfs);
cfs.invalidate(false);
@ -328,7 +329,8 @@ public class TrackerTest
Assert.assertEquals(0, tracker.getView().flushingMemtables.size());
Assert.assertEquals(0, cfs.metric.liveDiskSpaceUsed.getCount());
Assert.assertEquals(5, listener.received.size());
Assert.assertEquals(prev1, ((MemtableSwitchedNotification) listener.received.get(0)).memtable);
Assert.assertEquals(prev1, ((MemtableSwitchedNotification) listener.received.get(0)).previous);
Assert.assertEquals(next1, ((MemtableSwitchedNotification) listener.received.get(0)).next);
Assert.assertEquals(singleton(reader), ((SSTableAddedNotification) listener.received.get(1)).added);
Assert.assertEquals(Optional.of(prev1), ((SSTableAddedNotification) listener.received.get(1)).memtable());
Assert.assertEquals(prev1, ((MemtableDiscardedNotification) listener.received.get(2)).memtable);

10
test/unit/org/apache/cassandra/index/sai/SAITester.java
View File

@ -556,6 +556,11 @@ public abstract class SAITester extends CQLTester
}
protected void verifySSTableIndexes(IndexIdentifier indexIdentifier, int sstableContextCount, int sstableIndexCount)
{
verifySSTableIndexes(indexIdentifier, sstableContextCount, sstableIndexCount, 0);
}
protected void verifySSTableIndexes(IndexIdentifier indexIdentifier, int sstableContextCount, int sstableIndexCount, int expectedEmptyIndexCount)
{
ColumnFamilyStore cfs = getCurrentColumnFamilyStore();
StorageAttachedIndexGroup indexGroup = getCurrentIndexGroup();
@ -564,7 +569,10 @@ public abstract class SAITester extends CQLTester
StorageAttachedIndex sai = (StorageAttachedIndex) cfs.indexManager.getIndexByName(indexIdentifier.indexName);
Collection<SSTableIndex> sstableIndexes = sai == null ? Collections.emptyList() : sai.view().getIndexes();
assertEquals("Expected " + sstableIndexCount +" SSTableIndexes, but got " + sstableIndexes.toString(), sstableIndexCount, sstableIndexes.size());
long nonEmptyIndexCount = sstableIndexes.stream().filter(i -> i.getRowCount() > 0).count();
long emptyIndexCount = sstableIndexes.stream().filter(i -> i.getRowCount() == 0).count();
assertEquals("Expected " + sstableIndexCount +" SSTableIndexes, but got " + sstableIndexes.toString(), sstableIndexCount, nonEmptyIndexCount);
assertEquals("Expected " + expectedEmptyIndexCount + " empty indexes, but got " + emptyIndexCount, expectedEmptyIndexCount, emptyIndexCount);
}
protected boolean isBuildCompletionMarker(IndexComponent indexComponent)

8
test/unit/org/apache/cassandra/index/sai/cql/StorageAttachedIndexDDLTest.java
View File

@ -1167,8 +1167,8 @@ public class StorageAttachedIndexDDLTest extends SAITester
IndexTermType numericIndexTermType = createIndexTermType(Int32Type.instance);
IndexTermType literalIndexTermType = createIndexTermType(UTF8Type.instance);
populateData.run();
verifySSTableIndexes(numericIndexIdentifier, 2, 0);
verifySSTableIndexes(literalIndexIdentifier, 2, 0);
verifySSTableIndexes(numericIndexIdentifier, 2, 0, 2);
verifySSTableIndexes(literalIndexIdentifier, 2, 0, 2);
verifyIndexFiles(numericIndexTermType, numericIndexIdentifier, 2, 0, 2);
verifyIndexFiles(literalIndexTermType, literalIndexIdentifier, 2, 0, 2);
@ -1179,8 +1179,8 @@ public class StorageAttachedIndexDDLTest extends SAITester
// compact empty index
compact();
verifySSTableIndexes(numericIndexIdentifier, 1, 0);
verifySSTableIndexes(literalIndexIdentifier, 1, 0);
verifySSTableIndexes(numericIndexIdentifier, 1, 0, 1);
verifySSTableIndexes(literalIndexIdentifier, 1, 0, 1);
waitForAssert(() -> verifyIndexFiles(numericIndexTermType, numericIndexIdentifier, 1, 0, 1));
waitForAssert(() -> verifyIndexFiles(literalIndexTermType, literalIndexIdentifier, 1, 0, 1));

143
test/unit/org/apache/cassandra/index/sai/cql/VectorSiftSmallTest.java
View File

@ -34,10 +34,11 @@ import java.util.stream.IntStream;
import org.junit.Test;
import org.apache.cassandra.cql3.UntypedResultSet;
import org.apache.cassandra.index.sai.SAITester;
import static org.junit.Assert.assertTrue;
public class VectorSiftSmallTest extends VectorTester
public class VectorSiftSmallTest extends SAITester
{
@Test
public void testSiftSmall() throws Throwable
@ -60,6 +61,60 @@ public class VectorSiftSmallTest extends VectorTester
assertTrue("Disk recall is " + diskRecall, diskRecall > 0.95);
}
@Test
public void testSiftSmallWithBooleanPredicatesOfVaryingSelectivity() throws Throwable
{
var siftName = "siftsmall";
var baseVectors = readFvecs(String.format("test/data/%s/%s_base.fvecs", siftName, siftName));
var queryVectors = readFvecs(String.format("test/data/%s/%s_query.fvecs", siftName, siftName));
// Create table with regular id column and add SAI index on it
createTable("CREATE TABLE %s (pk int, id int, val vector<float, 128>, PRIMARY KEY(pk))");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
createIndex("CREATE CUSTOM INDEX ON %s(id) USING 'StorageAttachedIndex'");
// Insert all vectors into unique partitions with id column
insertVectorsWithId(baseVectors);
int totalVectors = baseVectors.size();
int topK = 100;
// Test with tiny range restriction (1% of data)
int tinyRangeEnd = totalVectors / 100;
// Test with small range restriction (10% of data)
int smallRangeEnd = totalVectors / 10;
// Test with medium range restriction (50% of data)
int mediumRangeEnd = totalVectors / 2;
// Execute queries of different selectivity. We compute the ground truth for each query vector by brute
// force in this test class because the dataset doesn't exist.
beforeAndAfterFlush(() -> {
double tinyRangeRecall = testRecallWithIdRange(queryVectors, baseVectors, 0, tinyRangeEnd, topK);
assertTrue("Tiny range recall is " + tinyRangeRecall, tinyRangeRecall >= 0.99);
double smallRangeRecall = testRecallWithIdRange(queryVectors, baseVectors, 0, smallRangeEnd, topK);
assertTrue("Small range recall is " + smallRangeRecall, smallRangeRecall >= 0.975);
double mediumRangeRecall = testRecallWithIdRange(queryVectors, baseVectors, 0, mediumRangeEnd, topK);
assertTrue("Medium range recall is " + mediumRangeRecall, mediumRangeRecall >= 0.975);
});
// Finish by deleting all rows and verifying queries return correctly.
// Using this test because it has many vectors already
disableCompaction();
for (int i = 0; i < totalVectors; i++)
execute("DELETE FROM %s WHERE pk = ?", i);
float[] vec = baseVectors.get(0);
beforeAndAfterFlush(() -> {
// Confirm all queries produce 0 rows. These queries hit several edge cases that are otherwise hard to
// test, so we add them at the end of this hybrid sift test.
assertRows(execute("SELECT pk FROM %s WHERE id >= ? AND id <= ? ORDER BY val ANN OF ? LIMIT ?", 0, tinyRangeEnd, vector(vec), 10));
assertRows(execute("SELECT pk FROM %s WHERE id >= ? AND id <= ? ORDER BY val ANN OF ? LIMIT ?", 0, smallRangeEnd, vector(vec), 10));
assertRows(execute("SELECT pk FROM %s WHERE id >= ? AND id <= ? ORDER BY val ANN OF ? LIMIT ?", 0, mediumRangeEnd, vector(vec), 10));
});
}
public static ArrayList<float[]> readFvecs(String filePath) throws IOException
{
var vectors = new ArrayList<float[]>();
@ -128,7 +183,7 @@ public class VectorSiftSmallTest extends VectorTester
UntypedResultSet result = execute("SELECT pk FROM %s ORDER BY val ANN OF " + queryVectorAsString + " LIMIT " + topK);
var gt = groundTruth.get(i);
int n = (int)result.stream().filter(row -> gt.contains(row.getInt("pk"))).count();
int n = (int) result.stream().filter(row -> gt.contains(row.getInt("pk"))).count();
topKfound.addAndGet(n);
}
catch (Throwable throwable)
@ -155,4 +210,88 @@ public class VectorSiftSmallTest extends VectorTester
}
});
}
private void insertVectorsWithId(List<float[]> baseVectors)
{
IntStream.range(0, baseVectors.size()).parallel().forEach(i -> {
float[] arrayVector = baseVectors.get(i);
String vectorAsString = Arrays.toString(arrayVector);
try
{
execute("INSERT INTO %s " + String.format("(pk, id, val) VALUES (%d, %d, %s)", i, i, vectorAsString));
}
catch (Throwable throwable)
{
throw new RuntimeException(throwable);
}
});
}
private double testRecallWithIdRange(List<float[]> queryVectors, List<float[]> baseVectors,
int idStart, int idEnd, int topK)
{
AtomicInteger topKfound = new AtomicInteger(0);
// Perform query with id range restriction and compute recall
IntStream.range(0, queryVectors.size()).parallel().forEach(i -> {
float[] queryVector = queryVectors.get(i);
String queryVectorAsString = Arrays.toString(queryVector);
try
{
// Compute ground truth for this filtered range by brute force
var filteredGroundTruth = computeGroundTruthForRange(queryVector, baseVectors, idStart, idEnd, topK);
String query = String.format("SELECT pk FROM %%s WHERE id >= %d AND id <= %d ORDER BY val ANN OF %s LIMIT %d",
idStart, idEnd, queryVectorAsString, topK);
UntypedResultSet result = execute(query);
// Count how many results are in the ground truth
int n = (int) result.stream().filter(row -> filteredGroundTruth.contains(row.getInt("pk"))).count();
topKfound.addAndGet(n);
}
catch (Throwable throwable)
{
throw new RuntimeException(throwable);
}
});
return (double) topKfound.get() / (queryVectors.size() * topK);
}
private HashSet<Integer> computeGroundTruthForRange(float[] queryVector, List<float[]> baseVectors,
int idStart, int idEnd, int topK)
{
// Create a list of (id, distance) pairs for vectors in the range
var candidates = new ArrayList<java.util.Map.Entry<Integer, Float>>();
for (int id = idStart; id <= idEnd && id < baseVectors.size(); id++)
{
float distance = euclideanDistance(queryVector, baseVectors.get(id));
candidates.add(new java.util.AbstractMap.SimpleEntry<>(id, distance));
}
// Sort by distance and take top K
candidates.sort(java.util.Map.Entry.comparingByValue());
var groundTruth = new HashSet<Integer>();
int limit = Math.min(topK, candidates.size());
for (int i = 0; i < limit; i++)
{
groundTruth.add(candidates.get(i).getKey());
}
return groundTruth;
}
private float euclideanDistance(float[] a, float[] b)
{
float sum = 0.0f;
for (int i = 0; i < a.length; i++)
{
float diff = a[i] - b[i];
sum += diff * diff;
}
return (float) Math.sqrt(sum);
}
}

39
test/unit/org/apache/cassandra/index/sai/cql/VectorTester.java
View File

@ -29,14 +29,17 @@ import io.github.jbellis.jvector.graph.GraphSearcher;
import io.github.jbellis.jvector.vector.VectorEncoding;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
import org.apache.cassandra.index.sai.SAITester;
import org.apache.cassandra.index.sai.disk.v1.segment.VectorIndexSegmentSearcher;
import org.apache.cassandra.index.sai.disk.v1.vector.ConcurrentVectorValues;
import org.apache.cassandra.index.sai.utils.Glove;
import org.apache.cassandra.inject.ActionBuilder;
import org.apache.cassandra.inject.Injections;
import org.apache.cassandra.inject.InvokePointBuilder;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Ignore;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
@Ignore
@RunWith(Parameterized.class)
public class VectorTester extends SAITester
{
protected static Glove.WordVector word2vec;
@ -47,29 +50,19 @@ public class VectorTester extends SAITester
word2vec = Glove.parse(VectorTester.class.getClassLoader().getResourceAsStream("glove.3K.50d.txt"));
}
@Parameterized.Parameter
public Boolean forceBruteForceQueries;
@Parameterized.Parameters(name = "forceBruteForceQueries={0}")
public static Iterable<Object[]> data()
{
return Arrays.asList(new Object[][]{{true}, {false}, {null}});
}
@Before
public void setup() throws Throwable
{
// override maxBruteForceRows to a random number between 0 and 4 so that we make sure
// the non-brute-force path gets called during tests (which mostly involve small numbers of rows)
var n = getRandom().nextIntBetween(0, 4);
var limitToTopResults = InvokePointBuilder.newInvokePoint()
.onClass("org.apache.cassandra.index.sai.disk.v2.V2VectorIndexSearcher")
.onMethod("limitToTopResults")
.atEntry();
var bitsOrPostingListForKeyRange = InvokePointBuilder.newInvokePoint()
.onClass("org.apache.cassandra.index.sai.disk.v2.V2VectorIndexSearcher")
.onMethod("bitsOrPostingListForKeyRange")
.atEntry();
var ab = ActionBuilder.newActionBuilder()
.actions()
.doAction("$this.globalBruteForceRows = " + n);
var changeBruteForceThreshold = Injections.newCustom("force_non_bruteforce_queries")
.add(limitToTopResults)
.add(bitsOrPostingListForKeyRange)
.add(ab)
.build();
Injections.inject(changeBruteForceThreshold);
VectorIndexSegmentSearcher.FORCE_BRUTE_FORCE_ANN = forceBruteForceQueries;
}
public static double rawIndexedRecall(Collection<float[]> vectors, float[] query, List<float[]> result, int topK) throws IOException

45
test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java
View File

@ -35,6 +35,7 @@ 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.index.sai.plan.QueryController;
import org.apache.cassandra.service.ClientWarn;
import static org.assertj.core.api.Assertions.assertThat;
@ -675,4 +676,48 @@ public class VectorTypeTest extends VectorTester
execute("INSERT INTO %s (pk, metadata, row_v) VALUES (10, {'map_k' : 'map_v'}, [0.11, 0.19])");
assertRows(execute(select), row);
}
@Test
public void testStaticVectorColumnIndex() throws Throwable
{
createTable("CREATE TABLE %s (pk int, ck int, val vector<float, 2> static, PRIMARY KEY(pk, ck))");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
execute("INSERT INTO %s (pk, ck, val) VALUES (0, 1, [1,0])");
execute("INSERT INTO %s (pk, ck) VALUES (0, 2)");
execute("INSERT INTO %s (pk, ck, val) VALUES (1, 3, [0,-1])");
execute("INSERT INTO %s (pk, ck, val) VALUES (2, 4, [0,1])");
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT ck FROM %s ORDER BY val ANN OF [0,1] LIMIT 3"), row(4), row(1), row(2));
assertRows(execute("SELECT ck FROM %s ORDER BY val ANN OF [0,1] LIMIT 2"), row(4), row(1));
});
}
@Test
public void testTooManyMaterializedKeys() throws Throwable
{
int originalValue = QueryController.MAX_MATERIALIZED_KEYS;
QueryController.MAX_MATERIALIZED_KEYS = 10;
try
{
createTable("CREATE TABLE %s (pk int primary key, i int, val vector<float, 2>)");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
createIndex("CREATE CUSTOM INDEX ON %s(i) USING 'StorageAttachedIndex'");
for (int i = 1; i <= QueryController.MAX_MATERIALIZED_KEYS * 10; i++)
execute("INSERT INTO %s (pk, i, val) VALUES (?, ?, [1,0])", i, i);
beforeAndAfterFlush(() -> {
// Search for less than half of the table, which is over the MAX_MATERIALIZED_KEYS value to trigger
// the switched order by then filter query execution.
UntypedResultSet rows = execute("SELECT pk FROM %s WHERE i < ? ORDER BY val ANN OF [0,1] LIMIT 3", QueryController.MAX_MATERIALIZED_KEYS * 2);
assertRowCount(rows, 3);
});
}
finally
{
QueryController.MAX_MATERIALIZED_KEYS = originalValue;
}
}
}

458
test/unit/org/apache/cassandra/index/sai/cql/VectorUpdateDeleteTest.java
View File

@ -19,11 +19,12 @@
package org.apache.cassandra.index.sai.cql;
import org.apache.cassandra.cql3.UntypedResultSet;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
import org.junit.Test;
import static org.apache.cassandra.config.CassandraRelevantProperties.SAI_VECTOR_SEARCH_ORDER_CHUNK_SIZE;
import static org.apache.cassandra.index.sai.cql.VectorTypeTest.assertContainsInt;
import static org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph.MIN_PQ_ROWS;
import static org.assertj.core.api.Assertions.assertThat;
public class VectorUpdateDeleteTest extends VectorTester
@ -82,6 +83,21 @@ public class VectorUpdateDeleteTest extends VectorTester
assertContainsInt(result, "pk", 0);
}
@Test
public void testFlushWithDeletedVectors()
{
createTable("CREATE TABLE %s (pk int, v vector<float, 2>, PRIMARY KEY(pk))");
createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
execute("INSERT INTO %s (pk, v) VALUES (0, [1.0, 2.0])");
execute("INSERT INTO %s (pk, v) VALUES (0, null)");
flush();
UntypedResultSet result = execute("SELECT * FROM %s ORDER BY v ann of [2.5, 3.5] LIMIT 1");
assertThat(result).hasSize(0);
}
// range delete won't trigger UpdateTransaction#onUpdated
@Test
public void rangeDeleteVectorInMemoryAndFlushTest()
@ -354,7 +370,7 @@ public class VectorUpdateDeleteTest extends VectorTester
execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
execute("UPDATE %s SET str_val='C' WHERE pk=0");
var result = execute("SELECT * FROM %s ORDER BY val ann of [0.5, 1.5, 2.5] LIMIT 2");
UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [0.5, 1.5, 2.5] LIMIT 2");
assertThat(result).hasSize(2);
}
@ -387,7 +403,7 @@ public class VectorUpdateDeleteTest extends VectorTester
execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'B', [2.0, 3.0, 4.0])");
flush();
var result = execute("SELECT * FROM %s ORDER BY val ann of [9.5, 10.5, 11.5] LIMIT 1");
UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [9.5, 10.5, 11.5] LIMIT 1");
assertThat(result).hasSize(1);
assertContainsInt(result, "pk", 0);
result = execute("SELECT * FROM %s ORDER BY val ann of [0.5, 1.5, 2.5] LIMIT 1");
@ -415,10 +431,36 @@ public class VectorUpdateDeleteTest extends VectorTester
flush();
// the shadow vector has the highest score
var result = execute("SELECT * FROM %s ORDER BY val ann of [1.0, 2.0, 3.0] LIMIT 1");
UntypedResultSet result = execute("SELECT * FROM %s ORDER BY val ann of [1.0, 2.0, 3.0] LIMIT 1");
assertThat(result).hasSize(1);
}
@Test
public void shadowedPrimaryKeyWithSharedVectorAndOtherPredicates()
{
createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, str_val text, val vector<float, 3>)");
createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
disableCompaction(KEYSPACE);
// flush a sstable with one vector that is shared by two rows
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 (2, 'A', [1.0, 2.0, 3.0])");
flush();
// flush another sstable to shadow row 0
execute("DELETE FROM %s where pk = 0");
flush();
// flush another sstable with one new vector row
execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'A', [2.0, 3.0, 4.0])");
flush();
// the shadowed vector has the highest score, but we shouldn't see it
UntypedResultSet result = execute("SELECT pk FROM %s WHERE str_val = 'A' ORDER BY val ann of [1.0, 2.0, 3.0] LIMIT 2");
assertRowsIgnoringOrder(result, row(2), row(1));
}
@Test
public void testVectorRowWhereUpdateMakesRowMatchNonOrderingPredicates()
{
@ -496,52 +538,386 @@ public class VectorUpdateDeleteTest extends VectorTester
assertRows(execute("SELECT pk FROM %s WHERE val1 = 'match me' AND val2 = 'match me' ORDER BY vec ANN OF [11,11] LIMIT 2"), row(1), row(2));
}
@Test
public void ensureVariableChunkSizeDoesNotLeadToIncorrectResults() throws Exception
public void shadowedPrimaryKeyWithUpdatedPredicateMatchingIntValue() throws Throwable
{
// When adding the chunk size feature, there were issues related to leaked files.
// This setting only matters for hybrid queries
createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, str_val text, vec vector<float, 2>)");
createIndex("CREATE CUSTOM INDEX ON %s(vec) USING 'StorageAttachedIndex' WITH OPTIONS = { 'similarity_function' : 'euclidean' }");
createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, num int, val vector<float, 3>)");
createIndex("CREATE CUSTOM INDEX ON %s(num) USING 'StorageAttachedIndex'");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
disableCompaction(KEYSPACE);
// Same PK, different num, different vectors
execute("INSERT INTO %s (pk, num, val) VALUES (0, 1, [1.0, 2.0, 3.0])");
execute("INSERT INTO %s (pk, num, val) VALUES (0, 2, [2.0, 2.0, 3.0])");
execute("INSERT INTO %s (pk, num, val) VALUES (0, 3, [3.0, 2.0, 3.0])");
// Need PKs that wrap 0 when put in PK order
execute("INSERT INTO %s (pk, num, val) VALUES (1, 1, [1.0, 2.0, 3.0])");
execute("INSERT INTO %s (pk, num, val) VALUES (2, 1, [1.0, 2.0, 3.0])");
// the shadowed vector has the highest score, but we shouldn't see it
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT pk FROM %s WHERE num < 3 ORDER BY val ann of [1.0, 2.0, 3.0] LIMIT 10"),
row(1), row(2));
});
}
@Test
public void rangeRestrictedTestWithDuplicateVectorsAndADelete()
{
createTable(String.format("CREATE TABLE %%s (pk int, str_val text, val vector<float, %d>, PRIMARY KEY(pk))", 2));
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
execute("INSERT INTO %s (pk, val) VALUES (0, [1.0, 2.0])"); // -3485513579396041028
execute("INSERT INTO %s (pk, val) VALUES (1, [1.0, 2.0])"); // -4069959284402364209
execute("INSERT INTO %s (pk, val) VALUES (2, [1.0, 2.0])"); // -3248873570005575792
execute("INSERT INTO %s (pk, val) VALUES (3, [1.0, 2.0])"); // 9010454139840013625
flush();
// Show the result set is as expected
assertRows(execute("SELECT pk FROM %s WHERE token(pk) <= -3248873570005575792 AND " +
"token(pk) >= -3485513579396041028 ORDER BY val ann of [1,2] LIMIT 1000"), row(0), row(2));
// Delete one of the rows
execute("DELETE FROM %s WHERE pk = 0");
flush();
assertRows(execute("SELECT pk FROM %s WHERE token(pk) <= -3248873570005575792 AND " +
"token(pk) >= -3485513579396041028 ORDER BY val ann of [1,2] LIMIT 1000"), row(2));
}
@Test
public void rangeRestrictedTestWithDuplicateVectorsAndAddNullVector() throws Throwable
{
createTable(String.format("CREATE TABLE %%s (pk int, str_val text, val vector<float, %d>, PRIMARY KEY(pk))", 2));
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
execute("INSERT INTO %s (pk, val) VALUES (0, [1.0, 2.0])");
execute("INSERT INTO %s (pk, val) VALUES (1, [1.0, 2.0])");
execute("INSERT INTO %s (pk, val) VALUES (2, [1.0, 2.0])");
// Add a str_val to make sure pk has a row id in the sstable
execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'a', null)");
// Add another row to test a different part of the code
execute("INSERT INTO %s (pk, val) VALUES (4, [1.0, 2.0])");
execute("DELETE FROM %s WHERE pk = 2");
flush();
// Delete one of the rows to trigger a shadowed primary key
execute("DELETE FROM %s WHERE pk = 0");
execute("INSERT INTO %s (pk, val) VALUES (2, [2.0, 2.0])");
flush();
// Delete more rows.
execute("DELETE FROM %s WHERE pk = 2");
execute("DELETE FROM %s WHERE pk = 3");
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1,2] LIMIT 1000"),
row(1), row(4));
});
}
// This test intentionally has extra rows with primary keys that are above and below the
// deleted primary key so that we do not short circuit certain parts of the shadowed key logic.
@Test
public void shadowedPrimaryKeyInDifferentSSTableEachWithMultipleRows()
{
createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, str_val text, val vector<float, 3>)");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
disableCompaction(KEYSPACE);
// flush a sstable with one vector
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])");
execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'A', [1.0, 2.0, 3.0])");
flush();
// flush another sstable to shadow the vector row
execute("INSERT INTO %s (pk, str_val, val) VALUES (1, 'A', [1.0, 2.0, 3.0])");
execute("DELETE FROM %s where pk = 2");
execute("INSERT INTO %s (pk, str_val, val) VALUES (3, 'A', [1.0, 2.0, 3.0])");
flush();
// flush another sstable with one new vector row
execute("INSERT INTO %s (pk, str_val, val) VALUES (0, 'B', [2.0, 3.0, 4.0])");
execute("INSERT INTO %s (pk, str_val, val) VALUES (4, 'B', [2.0, 3.0, 4.0])");
flush();
// the shadow vector has the highest score
UntypedResultSet result = execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 2.0, 3.0] LIMIT 4");
assertRows(result, row(1), row(3), row(0), row(4));
}
@Test
public void shadowedPrimaryKeysRequireDeeperSearch() throws Throwable
{
createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, str_val text, val vector<float, 2>)");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
disableCompaction(KEYSPACE);
// Create many sstables to ensure chunk size matters
// Start at 1 to prevent indexing zero vector.
// Index every vector with A to match everything and because this test only makes sense for hybrid queries
for (int i = 1; i <= 100; i++)
// Choose a row count that will essentially force us to re-query the index that still has more rows to search.
int baseRowCount = 1000;
// Create 1000 rows so that each row has a slightly less similar score.
for (int i = 0; i < baseRowCount - 10; i++)
{
execute("INSERT INTO %s (pk, str_val, vec) VALUES (?, ?, ?)", i, "A", vector((float) i, (float) i));
if (i % 10 == 0)
flush();
// Add some deletes in the next segment
if (i % 3 == 0)
execute("DELETE FROM %s WHERE pk = ?", i);
}
try
{
// We use a chunk size that is as low as possible (1) and goes up to the whole dataset (100).
// We also query for different LIMITs
for (int i = 1; i <= 100; i++)
try
{
SAI_VECTOR_SEARCH_ORDER_CHUNK_SIZE.setInt(i);
var results = execute("SELECT pk FROM %s WHERE str_val = 'A' ORDER BY vec ANN OF [1,1] LIMIT 1");
assertRows(results, row(1));
results = execute("SELECT pk FROM %s WHERE str_val = 'A' ORDER BY vec ANN OF [1,1] LIMIT 3");
// Note that we delete row 3
assertRows(results, row(1), row(2), row(4));
results = execute("SELECT pk FROM %s WHERE str_val = 'A' ORDER BY vec ANN OF [1,1] LIMIT 10");
// Note that we delete row 3, 6, 9, 12
assertRows(results, row(1), row(2), row(4), row(5),
row(7), row(8), row(10), row(11), row(13), row(14));
execute("INSERT INTO %s (pk, str_val, val) VALUES (?, 'A', ?)", i, vector(1f, (float) i));
}
catch (Error e)
{
logger.error("Failed to insert row {}", i, e);
throw new RuntimeException(e);
}
}
finally
for (int i = baseRowCount -10; i < baseRowCount; i++)
execute("INSERT INTO %s (pk, str_val, val) VALUES (?, 'A', ?)", i, vector(1f, (float) -i));
flush();
// Create 10 rows with the worst scores, but they won't be shadowed.
for (int i = baseRowCount; i < baseRowCount + 10; i++)
execute("INSERT INTO %s (pk, str_val, val) VALUES (?, 'A', ?)", i, vector(-1f, (float) baseRowCount * -1));
// Delete all but the last 10 rows
for (int i = 0; i < baseRowCount - 10; i++)
execute("DELETE FROM %s WHERE pk = ?", i);
beforeAndAfterFlush(() -> {
// ANN Only
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 1.0] LIMIT 3"),
row(baseRowCount - 10), row(baseRowCount - 9), row(baseRowCount - 8));
// Hyrbid
assertRows(execute("SELECT pk FROM %s WHERE str_val = 'A' ORDER BY val ann of [1.0, 1.0] LIMIT 3"),
row(baseRowCount - 10), row(baseRowCount - 9), row(baseRowCount - 8));
});
}
@Test
public void testUpdateVectorToWorseAndBetterPositions() throws Throwable
{
createTable("CREATE TABLE %s (pk int, val vector<float, 2>, PRIMARY KEY(pk))");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
execute("INSERT INTO %s (pk, val) VALUES (0, [1.0, 2.0])");
execute("INSERT INTO %s (pk, val) VALUES (1, [1.0, 3.0])");
flush();
execute("INSERT INTO %s (pk, val) VALUES (0, [1.0, 4.0])");
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 2.0] LIMIT 1"), row(1));
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 2.0] LIMIT 2"), row(1), row(0));
});
// And now update pk 1 to show that we can get 0 too
execute("INSERT INTO %s (pk, val) VALUES (1, [1.0, 5.0])");
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 2.0] LIMIT 1"), row(0));
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 2.0] LIMIT 2"), row(0), row(1));
});
// And now update both PKs so that the stream of ranked rows is PKs: 0, 1, [1], 0, 1, [0], where the numbers
// wrapped in brackets are the "real" scores of the vectors. This test makes sure that we correctly remove
// PrimaryKeys from the updatedKeys map so that we don't accidentally duplicate PKs.
execute("INSERT INTO %s (pk, val) VALUES (1, [1.0, 3.5])");
execute("INSERT INTO %s (pk, val) VALUES (0, [1.0, 6.0])");
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 2.0] LIMIT 1"), row(1));
assertRows(execute("SELECT pk FROM %s ORDER BY val ann of [1.0, 2.0] LIMIT 2"), row(1), row(0));
});
}
@Test
public void updatedPrimaryKeysRequireResumeSearch() throws Throwable
{
createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, str_val text, val vector<float, 2>)");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
disableCompaction(KEYSPACE);
// This test is fairly contrived, but it covers a bug we hit due to prematurely closed iterators.
// The general design for this test is to shadow the close vectors on a memtable/sstable index forcing the
// index to resume search. We do that by overwriting the first 50 vectors in the initial sstable.
for (int i = 0; i < 100; i++)
execute("INSERT INTO %s (pk, str_val, val) VALUES (?, 'A', ?)", i, vector(1f, (float) i));
// Add more rows to make sure we filter then sort
for (int i = 100; i < 1000; i++)
execute("INSERT INTO %s (pk, str_val, val) VALUES (?, 'C', ?)", i, vector(1f, (float) i));
flush();
// Overwrite the most similar 50 rows
for (int i = 0; i < 50; i++)
execute("INSERT INTO %s (pk, str_val, val) VALUES (?, 'B', ?)", i, vector(1f, (float) i));
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT pk FROM %s WHERE str_val = 'A' ORDER BY val ann of [1.0, 1.0] LIMIT 1"),
row(50));
});
}
@Test
public void testBruteForceRangeQueryWithUpdatedVectors1536D() throws Throwable
{
testBruteForceRangeQueryWithUpdatedVectors(1536);
}
@Test
public void testBruteForceRangeQueryWithUpdatedVectors2D() throws Throwable
{
testBruteForceRangeQueryWithUpdatedVectors(2);
}
private void testBruteForceRangeQueryWithUpdatedVectors(int vectorDimension) throws Throwable
{
createTable("CREATE TABLE %s (pk int, val vector<float, " + vectorDimension + ">, PRIMARY KEY(pk))");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
// Insert 100 vectors
for (int i = 0; i < 100; i++)
execute("INSERT INTO %s (pk, val) VALUES (?, ?)", i, randomVectorBoxed(vectorDimension));
// Update those vectors so some ordinals are changed
for (int i = 0; i < 100; i++)
execute("INSERT INTO %s (pk, val) VALUES (?, ?)", i, randomVectorBoxed(vectorDimension));
// Delete the first 50 PKs.
for (int i = 0; i < 50; i++)
execute("DELETE FROM %s WHERE pk = ?", i);
// All of the above inserts and deletes are performed on the same index to verify internal index behavior
// for both memtables and sstables.
beforeAndAfterFlush(() -> {
// Query for the first 10 vectors, we don't care which.
// Use a range query to hit the right brute force code path
UntypedResultSet results = execute("SELECT pk FROM %s WHERE token(pk) < 0 ORDER BY val ann of ? LIMIT 10",
randomVectorBoxed(vectorDimension));
assertThat(results).hasSize(10);
// Make sure we don't get any of the deleted PKs
assertThat(results).allSatisfy(row -> assertThat(row.getInt("pk")).isGreaterThanOrEqualTo(50));
});
}
@Test
public void testVectorIndexWithAllOrdinalsDeletedAndSomeViaRangeDeletion()
{
createTable(KEYSPACE, "CREATE TABLE %s (pk int, a int, str_val text, val vector<float, 3>, PRIMARY KEY(pk, a))");
createIndex("CREATE CUSTOM INDEX ON %s(str_val) USING 'StorageAttachedIndex'");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
disableCompaction(KEYSPACE);
// Insert two rows with different vectors to get different ordinals
execute("INSERT INTO %s (pk, a, str_val, val) VALUES (1, 1, 'A', [1.0, 2.0, 3.0])");
execute("INSERT INTO %s (pk, a, str_val, val) VALUES (2, 1, 'A', [1.0, 2.0, 4.0])");
// Range delete the first row
execute("DELETE FROM %s WHERE pk = 1");
// Specifically delete the vector column second to hit a different code path.
execute("DELETE FROM %s WHERE pk = 2 AND a = 1");
// Insert another row without a vector
execute("INSERT INTO %s (pk, a, str_val) VALUES (2, 1, 'A')");
flush();
assertRows(execute("SELECT PK FROM %s WHERE str_val = 'A' ORDER BY val ann of [1.0, 2.0, 3.0] LIMIT 1"));
}
@Test
public void ensureCompressedVectorsCanFlush()
{
createTable("CREATE TABLE %s (pk int, val vector<float, 4>, PRIMARY KEY(pk))");
IndexIdentifier indexIdentifier = createIndexIdentifier(createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'"));
// insert enough vectors for pq plus 1 because we need quantization and we're deleting a row
for (int i = 0; i < MIN_PQ_ROWS + 1; i++)
execute("INSERT INTO %s (pk, val) VALUES (?, ?)", i, vector(randomVector(4)));
// Delete a single vector to trigger the regression
execute("DELETE from %s WHERE pk = 0");
flush();
verifySSTableIndexes(indexIdentifier, 1);
}
// This test mimics having rf > 1.
@Test
public void testSameRowInMultipleSSTablesWithSameTimestamp() 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'");
// We don't want compaction preventing us from hitting the intended code path.
disableCompaction();
// This test is fairly contrived, but covers the case where the first row we attempt to materialize in the
// ScoreOrderedResultRetriever is shadowed by a row in a different sstable. And then, when we go to pull in
// the next row, we find that the PK is already pulled in, so we need to skip it.
execute("INSERT INTO %s (pk, ck, val) VALUES (0, 0, [1.0, 2.0, 3.0])");
execute("INSERT INTO %s (pk, ck, val) VALUES (0, 1, [1.0, 2.0, 3.0]) USING TIMESTAMP 1");
flush();
// Now, delete row pk=0, ck=0 so that we can test that the shadowed row is not returned and that we need
// to get the next row from the score ordered iterator.
execute("DELETE FROM %s WHERE pk = 0 AND ck = 0");
execute("INSERT INTO %s (pk, ck, val) VALUES (0, 1, [1.0, 2.0, 3.0]) USING TIMESTAMP 1");
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT ck FROM %s ORDER BY val ANN OF [1.0, 2.0, 3.0] LIMIT 2"), row(1));
});
}
@Test
public void testMemtableInsertSearchUpdateSearchHandling()
{
createTable("CREATE TABLE %s (id text PRIMARY KEY, embedding vector<float, 5>)");
createIndex("CREATE CUSTOM INDEX ON %s(embedding) USING 'StorageAttachedIndex' " +
"WITH OPTIONS = {'similarity_function': 'dot_product'}");
// Insert initial data
execute("INSERT INTO %s (id, embedding) VALUES ('row1', [0.1, 0.1, 0.1, 0.1, 0.1])");
execute("INSERT INTO %s (id, embedding) VALUES ('row2', [0.9, 0.9, 0.9, 0.9, 0.9])");
// Query 100 times to try to guarantee all graph searchers are initialized
for (int i = 0; i < 100; i++)
{
// Revert to prevent interference with other tests. Note that a decreased chunk size can impact
// whether we compute the topk with brute force because it determines how many vectors get sent to the
// vector index.
SAI_VECTOR_SEARCH_ORDER_CHUNK_SIZE.setInt(100000);
// Initial vector search
UntypedResultSet initialSearch = execute("SELECT * FROM %s ORDER BY embedding ANN OF [0.8, 0.8, 0.8, 0.8, 0.8] LIMIT 1");
assertThat(initialSearch).hasSize(1);
}
// Update one of the rows (this update wasn't observed due to state leaked between queries previously)
execute("UPDATE %s SET embedding = [0.7, 0.7, 0.7, 0.7, 0.7] WHERE id = 'row1'");
// Query 100 times to make sure it works as expected
for (int j = 0; j < 100; j++)
{
// Get all data to verify we have 2 rows
UntypedResultSet allData = execute("SELECT * FROM %s ORDER BY embedding ANN OF [0.8, 0.8, 0.8, 0.8, 0.8] LIMIT 1000");
assertThat(allData).hasSize(2);
}
}
@Test
public void testUpdatedVectorStaticVectorColumnIndex() throws Throwable
{
createTable("CREATE TABLE %s (pk int, ck int, val vector<float, 2> static, PRIMARY KEY(pk, ck))");
createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");
// This counts as an update because the indexed column is static and is therefore operated on at the partition
// level.
execute("INSERT INTO %s (pk, ck, val) VALUES (0, 1, [0,2])");
execute("INSERT INTO %s (pk, ck, val) VALUES (0, 2, [1,0])");
execute("INSERT INTO %s (pk, ck, val) VALUES (1, 3, [0,1])");
beforeAndAfterFlush(() -> {
assertRows(execute("SELECT ck FROM %s ORDER BY val ANN OF [1,0] LIMIT 2"), row(1), row(2));
});
}
}

8
test/unit/org/apache/cassandra/index/sai/disk/v1/InvertedIndexSearcherTest.java
View File

@ -45,6 +45,7 @@ import org.apache.cassandra.index.sai.disk.v1.trie.LiteralIndexWriter;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;
@ -66,6 +67,12 @@ public class InvertedIndexSearcherTest extends SAIRandomizedTester
{
private final PrimaryKey.Factory primaryKeyFactory = new PrimaryKey.Factory(Murmur3Partitioner.instance, new ClusteringComparator());
@Override
public SSTableId getSSTableId()
{
return null;
}
@Override
public PrimaryKey primaryKeyFromRowId(long sstableRowId)
{
@ -202,6 +209,7 @@ public class InvertedIndexSearcherTest extends SAIRandomizedTester
try (PerColumnIndexFiles indexFiles = new PerColumnIndexFiles(indexDescriptor, index.termType(), index.identifier()))
{
final IndexSegmentSearcher searcher = IndexSegmentSearcher.open(TEST_PRIMARY_KEY_MAP_FACTORY,
null,
indexFiles,
segmentMetadata,
index);

9
test/unit/org/apache/cassandra/index/sai/disk/v1/bbtree/BlockBalancedTreeIndexBuilder.java
View File

@ -52,6 +52,7 @@ import org.apache.cassandra.index.sai.memory.MemtableTermsIterator;
import org.apache.cassandra.index.sai.utils.IndexTermType;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.TermsIterator;
import org.apache.cassandra.io.sstable.SSTableId;
import org.apache.cassandra.utils.AbstractGuavaIterator;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;
@ -68,6 +69,12 @@ public class BlockBalancedTreeIndexBuilder
{
private final PrimaryKey.Factory primaryKeyFactory = new PrimaryKey.Factory(Murmur3Partitioner.instance, null);
@Override
public SSTableId getSSTableId()
{
return null;
}
@Override
public PrimaryKey primaryKeyFromRowId(long sstableRowId)
{
@ -140,7 +147,7 @@ public class BlockBalancedTreeIndexBuilder
try (PerColumnIndexFiles indexFiles = new PerColumnIndexFiles(indexDescriptor, index.termType(), index.identifier()))
{
IndexSegmentSearcher searcher = IndexSegmentSearcher.open(TEST_PRIMARY_KEY_MAP_FACTORY, indexFiles, metadata, index);
IndexSegmentSearcher searcher = IndexSegmentSearcher.open(TEST_PRIMARY_KEY_MAP_FACTORY, null, indexFiles, metadata, index);
assertThat(searcher, is(instanceOf(NumericIndexSegmentSearcher.class)));
return (NumericIndexSegmentSearcher) searcher;
}

4
test/unit/org/apache/cassandra/index/sai/functional/FlushingTest.java
View File

@ -69,13 +69,13 @@ public class FlushingTest extends SAITester
ResultSet rows = executeNet("SELECT id1 FROM %s WHERE v1 >= 0");
assertEquals(0, rows.all().size());
verifyIndexFiles(indexTermType, indexIdentifier, sstables, 0, sstables);
verifySSTableIndexes(indexIdentifier, sstables, 0);
verifySSTableIndexes(indexIdentifier, sstables, 0, 3);
compact();
waitForAssert(() -> verifyIndexFiles(indexTermType, indexIdentifier, 1, 0, 1));
rows = executeNet("SELECT id1 FROM %s WHERE v1 >= 0");
assertEquals(0, rows.all().size());
verifySSTableIndexes(indexIdentifier, 1, 0);
verifySSTableIndexes(indexIdentifier, 1, 0, 1);
}
}

3
test/unit/org/apache/cassandra/index/sai/functional/GroupComponentsTest.java
View File

@ -31,6 +31,7 @@ import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.index.sai.SAITester;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.StorageAttachedIndexGroup;
import org.apache.cassandra.index.sai.disk.EmptyIndex;
import org.apache.cassandra.index.sai.disk.format.Version;
import org.apache.cassandra.index.sai.utils.IndexTermType;
import org.apache.cassandra.io.sstable.Component;
@ -61,7 +62,7 @@ public class GroupComponentsTest extends SAITester
// index files are released but not removed
cfs.invalidate(true, false);
Assert.assertTrue(index.view().getIndexes().isEmpty());
Assert.assertTrue(index.view().getIndexes().stream().allMatch(i -> i instanceof EmptyIndex));
for (Component component : components)
Assert.assertTrue(sstable.descriptor.fileFor(component).exists());
}

55
test/unit/org/apache/cassandra/index/sai/memory/VectorMemoryIndexTest.java
View File

@ -29,10 +29,10 @@ import java.util.TreeMap;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import com.google.common.collect.Sets;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import org.mockito.Mockito;
import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.cql3.Operator;
@ -49,6 +49,7 @@ import org.apache.cassandra.db.filter.RowFilter;
import org.apache.cassandra.db.marshal.FloatType;
import org.apache.cassandra.db.marshal.Int32Type;
import org.apache.cassandra.db.marshal.VectorType;
import org.apache.cassandra.db.memtable.Memtable;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.dht.BootStrapper;
import org.apache.cassandra.dht.Bounds;
@ -59,19 +60,17 @@ import org.apache.cassandra.dht.Range;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.SAITester;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.RangeUtil;
import org.apache.cassandra.inject.Injections;
import org.apache.cassandra.inject.InvokePointBuilder;
import org.apache.cassandra.locator.TokenMetadata;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.FBUtilities;
import static org.apache.cassandra.config.CassandraRelevantProperties.MEMTABLE_SHARD_COUNT;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
@ -117,7 +116,9 @@ public class VectorMemoryIndexTest extends SAITester
@Test
public void randomQueryTest() throws Exception
{
memtableIndex = new VectorMemoryIndex(index);
// A non-null memtable tells it to track the mapping from primary key to vector, needed for brute force search
Memtable memtable = Mockito.mock(Memtable.class);
memtableIndex = new VectorMemoryIndex(index, memtable);
for (int row = 0; row < getRandom().nextIntBetween(1000, 5000); row++)
{
@ -130,6 +131,9 @@ public class VectorMemoryIndexTest extends SAITester
}
List<DecoratedKey> keys = new ArrayList<>(keyMap.keySet());
long actualVectorsReturned = 0;
long expectedVectorsReturned = 0;
double expectedRecall = 0.9;
for (int executionCount = 0; executionCount < 1000; executionCount++)
{
@ -149,28 +153,41 @@ public class VectorMemoryIndexTest extends SAITester
DataLimits.cqlLimits(limit),
DataRange.allData(cfs.metadata().partitioner));
try (KeyRangeIterator iterator = memtableIndex.search(new QueryContext(command,
DatabaseDescriptor.getRangeRpcTimeout(TimeUnit.MILLISECONDS)),
expression, keyRange))
long expectedResults = Math.min(limit, keysInRange.size());
try (CloseableIterator<PrimaryKeyWithScore> iterator = memtableIndex.orderBy(new QueryContext(command,
DatabaseDescriptor.getRangeRpcTimeout(TimeUnit.MILLISECONDS)),
expression, keyRange))
{
while (iterator.hasNext())
PrimaryKeyWithScore lastKey = null;
while (iterator.hasNext() && foundKeys.size() < expectedResults)
{
PrimaryKey primaryKey = iterator.next();
int key = Int32Type.instance.compose(primaryKey.partitionKey().getKey());
PrimaryKeyWithScore primaryKeyWithScore = iterator.next();
if (lastKey != null)
// This assertion only holds true as long as we query at most the expectedNumResults.
// Once we query deeper, we might get a key with a higher score than the last key.
// This is a direct consequence of the approximate part of ANN.
// Note that PrimaryKeyWithScore is flipped to descending order, so we use >= here.
assertTrue("Returned keys are not ordered by score", primaryKeyWithScore.compareTo(lastKey) >= 0);
lastKey = primaryKeyWithScore;
int key = Int32Type.instance.compose(primaryKeyWithScore.primaryKey().partitionKey().getKey());
assertFalse(foundKeys.contains(key));
assertTrue(keyRange.contains(primaryKey.partitionKey()));
assertTrue(keyRange.contains(primaryKeyWithScore.primaryKey().partitionKey()));
assertTrue(rowMap.containsKey(key));
foundKeys.add(key);
}
// Note that we weight each result evenly instead of each query evenly.
actualVectorsReturned += foundKeys.size();
expectedVectorsReturned += expectedResults;
if (foundKeys.size() < expectedResults)
assertTrue("Expected at least " + expectedResults + " results but got " + foundKeys.size(),
foundKeys.size() >= expectedResults * expectedRecall);
}
// with -Dcassandra.test.random.seed=260652334768666, there is one missing key
long expectedResult = Math.min(limit, keysInRange.size());
if (RangeUtil.coversFullRing(keyRange))
assertEquals("Missing key: " + Sets.difference(keysInRange, foundKeys), expectedResult, foundKeys.size());
else // if skip ANN, returned keys maybe larger than limit
assertTrue("Missing key: " + Sets.difference(keysInRange, foundKeys), expectedResult <= foundKeys.size());
}
assertTrue("Expected at least " + expectedVectorsReturned + " results but got " + actualVectorsReturned,
actualVectorsReturned >= expectedVectorsReturned * expectedRecall);
}
@Test