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.disk.v1.bbtree;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  
  import org.junit.Before;
  import org.junit.Test;
  
  import org.apache.cassandra.cql3.Operator;
  import org.apache.cassandra.db.marshal.Int32Type;
  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.IndexComponent;
  import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
  import org.apache.cassandra.index.sai.disk.v1.segment.SegmentTrieBuffer;
  import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
  import org.apache.cassandra.index.sai.metrics.QueryEventListener;
  import org.apache.cassandra.index.sai.plan.Expression;
  import org.apache.cassandra.index.sai.postings.PostingList;
  import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
  import org.apache.cassandra.io.util.FileHandle;
  import org.apache.cassandra.utils.FBUtilities;
  import org.apache.cassandra.utils.Throwables;
  import org.apache.lucene.index.PointValues.Relation;
  import org.apache.lucene.util.NumericUtils;
  
  import static org.apache.lucene.index.PointValues.Relation.CELL_CROSSES_QUERY;
  import static org.apache.lucene.index.PointValues.Relation.CELL_INSIDE_QUERY;
  import static org.apache.lucene.index.PointValues.Relation.CELL_OUTSIDE_QUERY;
  import static org.hamcrest.Matchers.greaterThan;
  import static org.hamcrest.Matchers.is;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertNotNull;
  import static org.junit.Assert.assertNull;
  import static org.junit.Assert.assertThat;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  public class BlockBalancedTreeReaderTest extends SAIRandomizedTester
  {
      private final BlockBalancedTreeReader.IntersectVisitor NONE_MATCH = new BlockBalancedTreeReader.IntersectVisitor()
      {
          @Override
          public boolean contains(byte[] packedValue)
          {
              return false;
          }
  
          @Override
          public Relation compare(byte[] minPackedValue, byte[] maxPackedValue)
          {
              return CELL_OUTSIDE_QUERY;
          }
      };
  
      private final BlockBalancedTreeReader.IntersectVisitor ALL_MATCH = new BlockBalancedTreeReader.IntersectVisitor()
      {
          @Override
          public boolean contains(byte[] packedValue)
          {
              return true;
          }
  
          @Override
          public Relation compare(byte[] minPackedValue, byte[] maxPackedValue)
          {
              return CELL_INSIDE_QUERY;
          }
      };
  
      private IndexDescriptor indexDescriptor;
      private StorageAttachedIndex index;
  
      @Before
      public void setup() throws Throwable
      {
          indexDescriptor = newIndexDescriptor();
          index = SAITester.createMockIndex(Int32Type.instance);
      }
  
      @Test
      public void testFilteringIntersection() throws Exception
      {
          int numRows = 1000;
  
          final SegmentTrieBuffer buffer = new SegmentTrieBuffer();
  
          for (int rowId = 0; rowId < numRows; rowId++)
          {
              buffer.add(integerToByteComparable(rowId), Integer.BYTES, rowId);
          }
  
          try (BlockBalancedTreeReader reader = finishAndOpenReader(4, buffer))
          {
              assertRange(reader, 445, 555);
          }
      }
  
      @Test
      public void testAdvance() throws Exception
      {
          final int numRows = between(1000, 2000);
          final SegmentTrieBuffer buffer = new SegmentTrieBuffer();
  
          for (int rowId = 0; rowId < numRows; rowId++)
          {
              buffer.add(integerToByteComparable(rowId), Integer.BYTES, rowId);
          }
  
          try (BlockBalancedTreeReader reader = finishAndOpenReader(2, buffer))
          {
              PostingList intersection = performIntersection(reader, NONE_MATCH);
              assertNull(intersection);
  
              intersection = performIntersection(reader, ALL_MATCH);
              assertEquals(numRows, intersection.size());
              assertEquals(100, intersection.advance(100));
              assertEquals(200, intersection.advance(200));
              assertEquals(300, intersection.advance(300));
              assertEquals(400, intersection.advance(400));
              assertEquals(401, intersection.advance(401));
              long expectedRowID = 402;
              for (long id = intersection.nextPosting(); expectedRowID < 500; id = intersection.nextPosting())
              {
                  assertEquals(expectedRowID++, id);
              }
              assertEquals(PostingList.END_OF_STREAM, intersection.advance(numRows + 1));
  
              intersection.close();
          }
      }
  
      @Test
      public void testSameValuesInLeaf() throws Exception
      {
          // While a bit synthetic this test is designed to test that the
          // BlockBalancedTreeReader.FilteringIntersection.buildPostingsFilterForSingleValueLeaf
          // method is exercised in a test. To do this we need to ensure that
          // we have at least one leaf that has all the same value and that
          // all of that leaf is requested in a query.
          final SegmentTrieBuffer buffer = new SegmentTrieBuffer();
  
          for (int rowId = 0; rowId < 10; rowId++)
          {
              buffer.add(integerToByteComparable(rowId), Integer.BYTES, rowId);
          }
  
          for (int rowId = 10; rowId < 20; rowId++)
          {
              buffer.add(integerToByteComparable(10), Integer.BYTES, rowId);
          }
  
          for (int rowId = 20; rowId < 30; rowId++)
          {
              buffer.add(integerToByteComparable(rowId), Integer.BYTES, rowId);
          }
  
          try (BlockBalancedTreeReader reader = finishAndOpenReader(5, buffer))
          {
              PostingList postingList = performIntersection(reader, buildQuery(8, 15));
  
              assertEquals(8, postingList.nextPosting());
              assertEquals(9, postingList.nextPosting());
              assertEquals(10, postingList.nextPosting());
              assertEquals(11, postingList.nextPosting());
              assertEquals(12, postingList.nextPosting());
              assertEquals(13, postingList.nextPosting());
              assertEquals(14, postingList.nextPosting());
              assertEquals(15, postingList.nextPosting());
              assertEquals(16, postingList.nextPosting());
              assertEquals(17, postingList.nextPosting());
              assertEquals(18, postingList.nextPosting());
              assertEquals(19, postingList.nextPosting());
              assertEquals(PostingList.END_OF_STREAM, postingList.nextPosting());
          }
      }
  
      @Test
      public void testResourcesReleaseWhenQueryDoesntMatchAnything() throws Exception
      {
          final SegmentTrieBuffer buffer = new SegmentTrieBuffer();
          for (int rowId = 0; rowId < 1000; rowId++)
          {
              buffer.add(integerToByteComparable(rowId), Integer.BYTES, rowId);
          }
          // add a gap between 1000 and 1100
          for (int rowId = 1000; rowId < 2000; rowId++)
          {
              buffer.add(integerToByteComparable(rowId + 100), Integer.BYTES, rowId);
          }
  
          try (BlockBalancedTreeReader reader = finishAndOpenReader(50, buffer))
          {
              final PostingList intersection = performIntersection(reader, buildQuery(1017, 1096));
              assertNull(intersection);
          }
      }
  
      @Test
      public void testConcurrentIntersectionsOnSameReader() throws Exception
      {
          int numRows = 1000;
  
          final SegmentTrieBuffer buffer = new SegmentTrieBuffer();
  
          for (int rowId = 0; rowId < numRows; rowId++)
          {
              buffer.add(integerToByteComparable(rowId), Integer.BYTES, rowId);
          }
  
          try (BlockBalancedTreeReader reader = finishAndOpenReader(4, buffer))
          {
              int concurrency = 100;
  
              ExecutorService executor = Executors.newFixedThreadPool(concurrency);
              List<Future<?>> results = new ArrayList<>();
              for (int thread = 0; thread < concurrency; thread++)
              {
                  results.add(executor.submit(() -> assertRange(reader, 445, 555)));
              }
              FBUtilities.waitOnFutures(results);
              executor.shutdown();
          }
      }
  
      @SuppressWarnings("SameParameterValue")
      private void assertRange(BlockBalancedTreeReader reader, long lowerBound, long upperBound)
      {
          Expression expression = Expression.create(index);
          expression.add(Operator.GT, Int32Type.instance.decompose(444));
          expression.add(Operator.LT, Int32Type.instance.decompose(555));
  
          try
          {
              PostingList intersection = performIntersection(reader, BlockBalancedTreeQueries.balancedTreeQueryFrom(expression, 4));
              assertNotNull(intersection);
              assertEquals(upperBound - lowerBound, intersection.size());
              for (long posting = lowerBound; posting < upperBound; posting++)
                  assertEquals(posting, intersection.nextPosting());
          }
          catch (IOException e)
          {
              throw Throwables.unchecked(e);
          }
      }
  
      private PostingList performIntersection(BlockBalancedTreeReader reader, BlockBalancedTreeReader.IntersectVisitor visitor)
      {
          QueryEventListener.BalancedTreeEventListener balancedTreeEventListener = mock(QueryEventListener.BalancedTreeEventListener.class);
          when(balancedTreeEventListener.postingListEventListener()).thenReturn(mock(QueryEventListener.PostingListEventListener.class));
          return reader.intersect(visitor, balancedTreeEventListener, mock(QueryContext.class));
      }
  
      private BlockBalancedTreeReader.IntersectVisitor buildQuery(int queryMin, int queryMax)
      {
          return new BlockBalancedTreeReader.IntersectVisitor()
          {
              @Override
              public boolean contains(byte[] packedValue)
              {
                  int x = NumericUtils.sortableBytesToInt(packedValue, 0);
                  return x >= queryMin && x <= queryMax;
              }
  
              @Override
              public Relation compare(byte[] minPackedValue, byte[] maxPackedValue)
              {
                  int min = NumericUtils.sortableBytesToInt(minPackedValue, 0);
                  int max = NumericUtils.sortableBytesToInt(maxPackedValue, 0);
                  assert max >= min;
  
                  if (max < queryMin || min > queryMax)
                  {
                      return Relation.CELL_OUTSIDE_QUERY;
                  }
                  else if (min >= queryMin && max <= queryMax)
                  {
                      return CELL_INSIDE_QUERY;
                  }
                  else
                  {
                      return CELL_CROSSES_QUERY;
                  }
              }
          };
      }
  
      private BlockBalancedTreeReader finishAndOpenReader(int maxPointsPerLeaf, SegmentTrieBuffer buffer) throws Exception
      {
          setBDKPostingsWriterSizing(8, 2);
          final NumericIndexWriter writer = new NumericIndexWriter(indexDescriptor,
                                                                   index.identifier(),
                                                                   maxPointsPerLeaf,
                                                                   Integer.BYTES);
  
          final SegmentMetadata.ComponentMetadataMap metadata = writer.writeCompleteSegment(buffer.iterator());
          final long treePosition = metadata.get(IndexComponent.BALANCED_TREE).root;
          assertThat(treePosition, is(greaterThan(0L)));
          final long postingsPosition = metadata.get(IndexComponent.POSTING_LISTS).root;
          assertThat(postingsPosition, is(greaterThan(0L)));
  
          FileHandle treeHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, index.identifier());
          FileHandle treePostingsHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, index.identifier());
          return new BlockBalancedTreeReader(index.identifier(),
                                             treeHandle,
                                             treePosition,
                                             treePostingsHandle,
                                             postingsPosition);
      }
  }
  