content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.cassandra.index.sai.memory;
  
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  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.apache.cassandra.ServerTestUtils;
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.cql3.Operator;
  import org.apache.cassandra.db.Clustering;
  import org.apache.cassandra.db.ColumnFamilyStore;
  import org.apache.cassandra.db.DataRange;
  import org.apache.cassandra.db.DecoratedKey;
  import org.apache.cassandra.db.PartitionPosition;
  import org.apache.cassandra.db.PartitionRangeReadCommand;
  import org.apache.cassandra.db.ReadCommand;
  import org.apache.cassandra.db.filter.ColumnFilter;
  import org.apache.cassandra.db.filter.DataLimits;
  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.dht.AbstractBounds;
  import org.apache.cassandra.dht.BootStrapper;
  import org.apache.cassandra.dht.Bounds;
  import org.apache.cassandra.dht.ExcludingBounds;
  import org.apache.cassandra.dht.IPartitioner;
  import org.apache.cassandra.dht.IncludingExcludingBounds;
  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.disk.format.Version;
  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.RangeUtil;
  import org.apache.cassandra.inject.Injections;
  import org.apache.cassandra.inject.InvokePointBuilder;
  import org.apache.cassandra.schema.TableMetadata;
  import org.apache.cassandra.tcm.ClusterMetadata;
  import org.apache.cassandra.utils.FBUtilities;
  
  import static org.apache.cassandra.config.CassandraRelevantProperties.MEMTABLE_SHARD_COUNT;
  import static org.apache.cassandra.config.CassandraRelevantProperties.ORG_APACHE_CASSANDRA_DISABLE_MBEAN_REGISTRATION;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertTrue;
  
  public class VectorMemoryIndexTest extends SAITester
  {
      private static final Injections.Counter indexSearchCounter = Injections.newCounter("IndexSearchCounter")
                                                                             .add(InvokePointBuilder.newInvokePoint()
                                                                                                    .onClass(TrieMemoryIndex.class)
                                                                                                    .onMethod("search"))
                                                                             .build();
  
      private ColumnFamilyStore cfs;
      private StorageAttachedIndex index;
      private VectorMemoryIndex memtableIndex;
      private IPartitioner partitioner;
      private Map<DecoratedKey, Integer> keyMap;
      private Map<Integer, ByteBuffer> rowMap;
      private int dimensionCount;
  
      @BeforeClass
      public static void setUpClass()
      {
          setupSeed();
          // Because this test wants to explicitly set tokens for the local node, we override SAITester::setUpClass, as
          // that calls CQLTester::setUpClass, which is opinionated about the locally owned tokens.
          MEMTABLE_SHARD_COUNT.setInt(8);
          ORG_APACHE_CASSANDRA_DISABLE_MBEAN_REGISTRATION.setBoolean(true);
  
          updateConfigs();
  
          ServerTestUtils.prepareServerNoRegister();
          //TODO (repeatability): this uses ThreadLocalRandom which doesn't let you control the seed, so this class can not be repeated...
          ServerTestUtils.registerLocal(BootStrapper.getRandomTokens(ClusterMetadata.current(), 10));
          ServerTestUtils.markCMS();
          // Ensure that the on-disk format statics are loaded before the test run
          Version.LATEST.onDiskFormat();
      }
  
      @Before
      public void setup() throws Throwable
      {
          dimensionCount = getRandom().nextIntBetween(2, 2048);
          index = SAITester.createMockIndex(VectorType.getInstance(FloatType.instance, dimensionCount));
          cfs = index.baseCfs();
          partitioner = cfs.getPartitioner();
          indexSearchCounter.reset();
          keyMap = new TreeMap<>();
          rowMap = new HashMap<>();
  
          Injections.inject(indexSearchCounter);
      }
  
      public static void reassignLocalTokens()
      {
      }
  
      @Test
      public void randomQueryTest() throws Exception
      {
          memtableIndex = new VectorMemoryIndex(index);
  
          for (int row = 0; row < getRandom().nextIntBetween(1000, 5000); row++)
          {
              int pk = getRandom().nextIntBetween(0, 10000);
              while (rowMap.containsKey(pk))
                  pk = getRandom().nextIntBetween(0, 10000);
              var value = randomVector();
              rowMap.put(pk, value);
              addRow(pk, value);
          }
  
          List<DecoratedKey> keys = new ArrayList<>(keyMap.keySet());
  
          for (int executionCount = 0; executionCount < 1000; executionCount++)
          {
              Expression expression = generateRandomExpression();
              AbstractBounds<PartitionPosition> keyRange = generateRandomBounds(keys);
              Set<Integer> keysInRange = keys.stream().filter(keyRange::contains)
                                             .map(k -> Int32Type.instance.compose(k.getKey()))
                                             .collect(Collectors.toSet());
  
              Set<Integer> foundKeys = new HashSet<>();
              int limit = getRandom().nextIntBetween(1, 100);
  
              ReadCommand command = PartitionRangeReadCommand.create(cfs.metadata(),
                                                                     FBUtilities.nowInSeconds(),
                                                                     ColumnFilter.all(cfs.metadata()),
                                                                     RowFilter.none(),
                                                                     DataLimits.cqlLimits(limit),
                                                                     DataRange.allData(cfs.metadata().partitioner));
  
              try (KeyRangeIterator iterator = memtableIndex.search(new QueryContext(command,
                                                                                     DatabaseDescriptor.getRangeRpcTimeout(TimeUnit.MILLISECONDS)),
                                                                    expression, keyRange))
              {
                  while (iterator.hasNext())
                  {
                      PrimaryKey primaryKey = iterator.next();
                      int key = Int32Type.instance.compose(primaryKey.partitionKey().getKey());
                      assertFalse(foundKeys.contains(key));
  
                      assertTrue(keyRange.contains(primaryKey.partitionKey()));
                      assertTrue(rowMap.containsKey(key));
                      foundKeys.add(key);
                  }
              }
              // 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());
          }
      }
  
      @Test
      public void indexIteratorTest()
      {
          // VSTODO
      }
  
      private Expression generateRandomExpression()
      {
          Expression expression = Expression.create(index);
          expression.add(Operator.ANN, randomVector());
          return expression;
      }
  
      private ByteBuffer randomVector() {
          List<Float> rawVector = new ArrayList<>(dimensionCount);
          for (int i = 0; i < dimensionCount; i++) {
              rawVector.add(getRandom().nextFloat());
          }
          return VectorType.getInstance(FloatType.instance, dimensionCount).getSerializer().serialize(rawVector);
      }
  
      private AbstractBounds<PartitionPosition> generateRandomBounds(List<DecoratedKey> keys)
      {
          PartitionPosition leftBound = getRandom().nextBoolean() ? partitioner.getMinimumToken().minKeyBound()
                                                                  : keys.get(getRandom().nextIntBetween(0, keys.size() - 1)).getToken().minKeyBound();
  
          PartitionPosition rightBound = getRandom().nextBoolean() ? partitioner.getMinimumToken().minKeyBound()
                                                                   : keys.get(getRandom().nextIntBetween(0, keys.size() - 1)).getToken().maxKeyBound();
  
          AbstractBounds<PartitionPosition> keyRange;
  
          if (leftBound.isMinimum() && rightBound.isMinimum())
              keyRange = new Range<>(leftBound, rightBound);
          else
          {
              if (AbstractBounds.strictlyWrapsAround(leftBound, rightBound))
              {
                  PartitionPosition temp = leftBound;
                  leftBound = rightBound;
                  rightBound = temp;
              }
              if (getRandom().nextBoolean())
                  keyRange = new Bounds<>(leftBound, rightBound);
              else if (getRandom().nextBoolean())
                  keyRange = new ExcludingBounds<>(leftBound, rightBound);
              else
                  keyRange = new IncludingExcludingBounds<>(leftBound, rightBound);
          }
          return keyRange;
      }
  
      private void addRow(int pk, ByteBuffer value)
      {
          DecoratedKey key = makeKey(cfs.metadata(), pk);
          memtableIndex.add(key, Clustering.EMPTY, value);
          keyMap.put(key, pk);
      }
  
      private DecoratedKey makeKey(TableMetadata table, Integer partitionKey)
      {
          ByteBuffer key = table.partitionKeyType.fromString(partitionKey.toString());
          return table.partitioner.decorateKey(key);
      }
  }
  