adelapena commented on code in PR #2673: URL: https://github.com/apache/cassandra/pull/2673#discussion_r1370280849
########## src/java/org/apache/cassandra/index/sai/plan/VectorTopKProcessor.java: ########## @@ -0,0 +1,198 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.cassandra.index.sai.plan; + +import java.nio.ByteBuffer; +import java.util.Comparator; +import java.util.Optional; +import java.util.PriorityQueue; +import java.util.TreeMap; +import java.util.TreeSet; +import javax.annotation.Nullable; + +import com.google.common.base.Preconditions; +import org.apache.commons.lang3.tuple.Triple; + +import org.apache.cassandra.cql3.Operator; +import org.apache.cassandra.db.ColumnFamilyStore; +import org.apache.cassandra.db.DecoratedKey; +import org.apache.cassandra.db.Keyspace; +import org.apache.cassandra.db.ReadCommand; +import org.apache.cassandra.db.filter.RowFilter; +import org.apache.cassandra.db.partitions.BasePartitionIterator; +import org.apache.cassandra.db.partitions.PartitionIterator; +import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator; +import org.apache.cassandra.db.rows.BaseRowIterator; +import org.apache.cassandra.db.rows.Row; +import org.apache.cassandra.db.rows.Unfiltered; +import org.apache.cassandra.index.Index; +import org.apache.cassandra.index.SecondaryIndexManager; +import org.apache.cassandra.index.sai.IndexContext; +import org.apache.cassandra.index.sai.StorageAttachedIndex; +import org.apache.cassandra.index.sai.utils.InMemoryPartitionIterator; +import org.apache.cassandra.index.sai.utils.InMemoryUnfilteredPartitionIterator; +import org.apache.cassandra.index.sai.utils.PartitionInfo; +import org.apache.cassandra.index.sai.utils.TypeUtil; +import org.apache.cassandra.schema.ColumnMetadata; +import org.apache.cassandra.utils.FBUtilities; +import org.apache.cassandra.utils.Pair; + +/** + * Processor that scans all rows from given partitions and selects rows with top-k scores based on vector indexes. + * <p> + * This processor performs the following steps: + * - collect rows with score into PriorityQueue that sorts rows based on score. If there are multiple vector indexes, Review Comment: ```suggestion * - collect rows with score into a {@link PriorityQueue} that sorts rows based on score. If there are multiple vector indexes, ``` ########## src/java/org/apache/cassandra/index/sai/iterators/KeyRangeListIterator.java: ########## @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.cassandra.index.sai.iterators; + +import java.util.List; + +import com.google.common.collect.Iterators; +import com.google.common.collect.PeekingIterator; + +import org.apache.cassandra.index.sai.utils.PrimaryKey; + +/** + * A {@link KeyRangeIterator} that iterates over a list of {@link PrimaryKey}s without modifying the underlying list. + */ +public class KeyRangeListIterator extends KeyRangeIterator +{ + private final PeekingIterator<PrimaryKey> keyQueue; + + /** + * Create a new {@link KeyRangeListIterator} that iterates over the provided list of keys. Review Comment: ```suggestion * Create a new {@link KeyRangeListIterator} that iterates over the provided list of keys. * ``` ########## src/java/org/apache/cassandra/index/sai/disk/v1/postings/VectorPostingList.java: ########## @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.cassandra.index.sai.disk.v1.postings; + +import java.io.IOException; +import java.util.PrimitiveIterator; + +import org.apache.cassandra.index.sai.postings.PostingList; +import org.apache.lucene.util.LongHeap; + +/** + * A posting list for ANN search results. Transforms result from similarity order to rowId order. Review Comment: ```suggestion * A {@link PostingList} for ANN search results. Transforms result from similarity order to row ID order. ``` ########## test/distributed/org/apache/cassandra/distributed/test/sai/VectorDistributedTest.java: ########## @@ -0,0 +1,440 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.cassandra.distributed.test.sai; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.Multimap; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; + +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import org.apache.cassandra.cql3.statements.SelectStatement; +import org.apache.cassandra.db.Keyspace; +import org.apache.cassandra.db.marshal.Int32Type; +import org.apache.cassandra.dht.Murmur3Partitioner; +import org.apache.cassandra.distributed.Cluster; +import org.apache.cassandra.distributed.api.ConsistencyLevel; +import org.apache.cassandra.distributed.test.TestBaseImpl; +import org.apache.cassandra.index.sai.SAITester; +import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig; + +import static org.apache.cassandra.distributed.api.Feature.GOSSIP; +import static org.apache.cassandra.distributed.api.Feature.NETWORK; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class VectorDistributedTest extends TestBaseImpl +{ + @Rule + public SAITester.FailureWatcher failureRule = new SAITester.FailureWatcher(); + + private static final String CREATE_KEYSPACE = "CREATE KEYSPACE %%s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': %d}"; + private static final String CREATE_TABLE = "CREATE TABLE %%s (pk int primary key, val vector<float, %d>)"; + private static final String CREATE_INDEX = "CREATE CUSTOM INDEX ON %%s(%s) USING 'StorageAttachedIndex'"; + + private static final VectorSimilarityFunction function = IndexWriterConfig.DEFAULT_SIMILARITY_FUNCTION; + + private static final double MIN_RECALL = 0.8; + + private static final int NUM_REPLICAS = 3; + private static final int RF = 2; + + private static final AtomicInteger seq = new AtomicInteger(); + private static String table; + + private static Cluster cluster; + + private static int dimensionCount; + + @BeforeClass + public static void setupCluster() throws Exception + { + cluster = Cluster.build(NUM_REPLICAS) + .withTokenCount(1) // VSTODO in-jvm-test in CC branch doesn't support multiple tokens Review Comment: We still have to get rid of the mention of Converged Cassandra. ########## src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java: ########## @@ -90,17 +99,35 @@ import org.apache.cassandra.schema.ColumnMetadata; import org.apache.cassandra.schema.IndexMetadata; import org.apache.cassandra.schema.TableMetadata; +import org.apache.cassandra.service.ClientWarn; import org.apache.cassandra.service.StorageService; import org.apache.cassandra.utils.Pair; import org.apache.cassandra.utils.concurrent.OpOrder; +import static org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig.MAX_TOP_K; + public class StorageAttachedIndex implements Index { public static final String NAME = "sai"; + + public static final String VECTOR_USAGE_WARNING = "SAI ANN indexes on vector columns are experimental and are not recommended for production use.\n" + Review Comment: I have just realized that the documentation for SAI committed by CASSANDRA-18231 a month ago contains documentation about vector search. We should verify that the doc matches this implementation, and probably add the list of query limitations to it, probably [here](https://github.com/apache/cassandra/blob/trunk/doc/modules/cassandra/partials/vector-search/query-vector-data.adoc). ########## src/java/org/apache/cassandra/index/sai/plan/VectorTopKProcessor.java: ########## @@ -0,0 +1,198 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.cassandra.index.sai.plan; + +import java.nio.ByteBuffer; +import java.util.Comparator; +import java.util.Optional; +import java.util.PriorityQueue; +import java.util.TreeMap; +import java.util.TreeSet; +import javax.annotation.Nullable; + +import com.google.common.base.Preconditions; +import org.apache.commons.lang3.tuple.Triple; + +import org.apache.cassandra.cql3.Operator; +import org.apache.cassandra.db.ColumnFamilyStore; +import org.apache.cassandra.db.DecoratedKey; +import org.apache.cassandra.db.Keyspace; +import org.apache.cassandra.db.ReadCommand; +import org.apache.cassandra.db.filter.RowFilter; +import org.apache.cassandra.db.partitions.BasePartitionIterator; +import org.apache.cassandra.db.partitions.PartitionIterator; +import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator; +import org.apache.cassandra.db.rows.BaseRowIterator; +import org.apache.cassandra.db.rows.Row; +import org.apache.cassandra.db.rows.Unfiltered; +import org.apache.cassandra.index.Index; +import org.apache.cassandra.index.SecondaryIndexManager; +import org.apache.cassandra.index.sai.IndexContext; +import org.apache.cassandra.index.sai.StorageAttachedIndex; +import org.apache.cassandra.index.sai.utils.InMemoryPartitionIterator; +import org.apache.cassandra.index.sai.utils.InMemoryUnfilteredPartitionIterator; +import org.apache.cassandra.index.sai.utils.PartitionInfo; +import org.apache.cassandra.index.sai.utils.TypeUtil; +import org.apache.cassandra.schema.ColumnMetadata; +import org.apache.cassandra.utils.FBUtilities; +import org.apache.cassandra.utils.Pair; + +/** + * Processor that scans all rows from given partitions and selects rows with top-k scores based on vector indexes. + * <p> + * This processor performs the following steps: + * - collect rows with score into PriorityQueue that sorts rows based on score. If there are multiple vector indexes, + * the final score is the sum of all vector index scores. + * - remove rows with the lowest scores from PQ if PQ size exceeds limit + * - return rows from PQ in primary key order to client + * <p> + * Note that recall will be lower with paging, because: + * - page size is used as limit + * - for the first query, coordinator returns global top page-size rows within entire ring + * - for the subsequent queries, coordinators returns global top page-size rows withom range from last-returned-row to max token Review Comment: We don't support paging, maybe this comes from a previous attempt to support it? ########## src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentOrdering.java: ########## @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.cassandra.index.sai.disk.v1.segment; + +import java.io.IOException; +import java.util.List; + +import org.apache.cassandra.dht.AbstractBounds; +import org.apache.cassandra.index.sai.QueryContext; +import org.apache.cassandra.index.sai.iterators.KeyRangeIterator; +import org.apache.cassandra.index.sai.plan.Expression; +import org.apache.cassandra.index.sai.utils.PrimaryKey; + +/** + * A {@link SegmentOrdering} orders and limits a list of {@link PrimaryKey}s. + * <p> + * When using {@link SegmentOrdering} there are several steps to + * build the list of Primary Keys to be ordered and limited: + * <p> + * 1. Find all primary keys that match each non-ordering query predicate. + * 2. Union and intersect the results of step 1 to build a single {@link KeyRangeIterator} + * ordered by {@link PrimaryKey}. + * 3. Filter out any shadowed primary keys. + * 4. Fan the primary keys from step 3 out to each sstable segment to order and limit each + * list of primary keys. + * <p> + * SegmentOrdering handles the fourth step. + * <p> + * Note: a segment ordering is only used when a query has both ordering and non-ordering predicates. + * Where a query has only ordering predicates, the ordering is handled by the Review Comment: ```suggestion * Where a query has only ordering predicates, the ordering is handled by ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
