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.iterators;
  
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.List;
  import java.util.stream.Collectors;
  
  import com.google.common.annotations.VisibleForTesting;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.cassandra.config.CassandraRelevantProperties;
  import org.apache.cassandra.index.sai.utils.PrimaryKey;
  import org.apache.cassandra.index.sai.utils.PrimaryKey.Kind;
  import org.apache.cassandra.io.util.FileUtils;
  import org.apache.cassandra.tracing.Tracing;
  
  import javax.annotation.Nullable;
  
  /**
   * A simple intersection iterator that makes no real attempts at optimising the iteration apart from
   * initially sorting the ranges. This implementation also supports an intersection limit via
   * {@code CassandraRelevantProperties.SAI_INTERSECTION_CLAUSE_LIMIT} which limits the number of ranges that will 
   * be included in the intersection. This currently defaults to 2.
   * <p> 
   * Intersection only works for ranges that are compatible according to {@link PrimaryKey.Kind#isIntersectable(Kind)}.
   */
  public class KeyRangeIntersectionIterator extends KeyRangeIterator
  {
      private static final Logger logger = LoggerFactory.getLogger(KeyRangeIntersectionIterator.class);
  
      static
      {
          logger.info(String.format("Storage attached index intersection clause limit is %d", CassandraRelevantProperties.SAI_INTERSECTION_CLAUSE_LIMIT.getInt()));
      }
  
      private final List<KeyRangeIterator> ranges;
      private PrimaryKey highestKey;
  
      private KeyRangeIntersectionIterator(Builder.Statistics statistics, List<KeyRangeIterator> ranges, Runnable onClose)
      {
          super(statistics, onClose);
          this.ranges = ranges;
          this.highestKey = null;
      }
  
      @Override
      protected PrimaryKey computeNext()
      {
          if (highestKey == null)
              highestKey = computeHighestKey();
  
          outer:
          // After advancing one iterator, we must try to advance all the other iterators that got behind,
          // so they catch up to it. Note that we will not advance the iterators for static columns
          // as long as they point to the partition of the highest key. (This is because STATIC primary keys 
          // compare to other keys only by partition.) This loop continues until all iterators point to the same key,
          // or if we run out of keys on any of them, or if we exceed the maximum key.
          // There is no point in iterating after maximum, because no keys will match beyond that point.
          while (highestKey != null && highestKey.compareTo(getMaximum(), false) <= 0)
          {
              // Try to advance all iterators to the highest key seen so far.
              // Once this inner loop finishes normally, all iterators are guaranteed to be at the same value.
              for (KeyRangeIterator range : ranges)
              {
                  if (!range.hasNext())
                      return endOfData();
  
                  if (range.peek().compareTo(highestKey, false) < 0)
                  {
                      PrimaryKey nextKey = skipToHighestKey(range);
  
                      // We use strict comparison here, since it orders WIDE primary keys after STATIC primary keys
                      // in the same partition. When WIDE keys are present, we want to return them rather than STATIC
                      // keys to avoid retrieving and post-filtering entire partitions.
                      if (nextKey == null || nextKey.compareTo(highestKey, true) > 0)
                      {
                          // We jumped over the highest key seen so far, so make it the new highest key.
                          highestKey = nextKey;
  
                          // This iterator jumped over, so the other iterators might be lagging behind now,
                          // including the ones already advanced in the earlier cycles of the inner loop.
                          // Therefore, restart the inner loop in order to advance the lagging iterators.
                          continue outer;
                      }
                      assert nextKey.compareTo(highestKey, false) == 0 :
                          String.format("Skipped to a key smaller than the target! " +
                                        "iterator: %s, target key: %s, returned key: %s", range, highestKey, nextKey);
                  }
              }
  
              // If we get here, all iterators have been advanced to the same key. When STATIC and WIDE keys are
              // mixed, this means WIDE keys point to exactly the same row, and STATIC keys the same partition.
              PrimaryKey result = highestKey;
  
              // Advance one iterator to the next key and remember the key as the highest seen so far.
              // It can become null when we reach the end of the iterator.
              // If there are both static and non-static keys being iterated here, we advance a non-static one,
              // regardless of the order of ranges in the ranges list.
              highestKey = advanceOneRange();
  
              // If we get here, all iterators have been advanced to the same key. When STATIC and WIDE keys are
              // mixed, this means WIDE keys point to exactly the same row, and STATIC keys the same partition.
              return result;
          }
  
          return endOfData();
      }
  
      private PrimaryKey skipToHighestKey(KeyRangeIterator range)
      {
          if (range.peek().kind() == highestKey.kind())
              return skipAndPeek(range, highestKey);
  
          if (range.peek().kind() == Kind.STATIC)
          {
              // If we advance a STATIC key, then we must advance it to the same partition as the highestKey.
              // Advancing a STATIC key to a WIDE key directly (without throwing away the clustering) would
              // go too far, as WIDE keys are stored after STATIC in the posting list.
              PrimaryKey nextKey = skipAndPeek(range, highestKey.toStatic());
  
              if (nextKey != null && nextKey.compareTo(highestKey, true) < 0 && nextKey.kind() == Kind.WIDE)
                  // This iterator may have mixed STATIC and non-STATIC postings. Advance again if we've 
                  // landed on a WIDE key that sorts lower in the same partition. 
                  nextKey = skipAndPeek(range, highestKey);
  
              return nextKey;
          }
  
          return skipAndPeek(range, highestKey);
      }
  
      /**
       * Advances the iterator of one range to the next item, which becomes the highest seen so far.
       * Iterators pointing to STATIC keys are advanced only if no non-STATIC keys have been advanced.
       *
       * @return the next highest key or null if the iterator has reached the end
       */
      private @Nullable PrimaryKey advanceOneRange()
      {
          for (KeyRangeIterator range : ranges)
              if (range.peek().kind() != Kind.STATIC)
              {
                  range.next();
                  return range.hasNext() ? range.peek() : null;
              }
          
          for (KeyRangeIterator range : ranges)
              if (range.peek().kind() == Kind.STATIC)
              {
                  range.next();
                  return range.hasNext() ? range.peek() : null;
              }
  
          throw new IllegalStateException("There should be at least one range to advance!");
      }
  
      private @Nullable PrimaryKey computeHighestKey()
      {
          PrimaryKey max = getMinimum();
          for (KeyRangeIterator range : ranges)
          {
              if (!range.hasNext())
                  return null;
              if (range.peek().compareTo(max, true) > 0)
                  max = range.peek();
          }
          return max;
      }
  
      @Override
      protected void performSkipTo(PrimaryKey nextKey)
      {
          // Resist the temptation to call range.hasNext before skipTo: this is a pessimisation, hasNext will invoke
          // computeNext under the hood, which is an expensive operation to produce a value that we plan to throw away.
          // Instead, it is the responsibility of the child iterators to make skipTo fast when the iterator is exhausted.
          for (KeyRangeIterator range : ranges)
              range.skipTo(nextKey);
  
          // Force recomputing the highest key on the next call to computeNext()
          highestKey = null;
      }
  
      @Override
      public void close()
      {
          super.close();
          FileUtils.closeQuietly(ranges);
      }
  
      /**
       * Fetches the next available item from the iterator, such that the item is not lower than the given key.
       * If no such items are available, returns null.
       */
      private PrimaryKey skipAndPeek(KeyRangeIterator iterator, PrimaryKey minKey)
      {
          iterator.skipTo(minKey);
          return iterator.hasNext() ? iterator.peek() : null;
      }
  
      public static Builder builder(int size, int limit)
      {
          return builder(size, limit, () -> {});
      }
  
      public static Builder builder(int size, Runnable onClose)
      {
          return new Builder(size, onClose);
      }
  
      @VisibleForTesting
      public static Builder builder(int size, int limit, Runnable onClose)
      {
          return new Builder(size, limit, onClose);
      }
  
      @VisibleForTesting
      public static class Builder extends KeyRangeIterator.Builder
      {
          // This controls the maximum number of range iterators that will be used in the final
          // intersection of a query operation. It is set from cassandra.sai.intersection_clause_limit
          // and defaults to 2
          private final int limit;
          // tracks if any of the added ranges are disjoint with the other ranges, which is useful
          // in case of intersection, as it gives a direct answer whether the iterator is going
          // to produce any results.
          private boolean isDisjoint;
  
          protected final List<KeyRangeIterator> rangeIterators;
  
          Builder(int size, Runnable onClose)
          {
              this(size, CassandraRelevantProperties.SAI_INTERSECTION_CLAUSE_LIMIT.getInt(), onClose);
          }
  
          Builder(int size, int limit, Runnable onClose)
          {
              super(new IntersectionStatistics(), onClose);
              rangeIterators = new ArrayList<>(size);
              this.limit = limit;
          }
  
          @Override
          public KeyRangeIterator.Builder add(KeyRangeIterator range)
          {
              if (range == null)
                  return this;
  
              if (range.getMaxKeys() > 0)
                  rangeIterators.add(range);
              else
                  FileUtils.closeQuietly(range);
  
              updateStatistics(statistics, range);
  
              return this;
          }
  
          @Override
          public int rangeCount()
          {
              return rangeIterators.size();
          }
  
          @Override
          public void cleanup()
          {
              super.cleanup();
              FileUtils.closeQuietly(rangeIterators);
          }
  
          @Override
          protected KeyRangeIterator buildIterator()
          {
              rangeIterators.sort(Comparator.comparingLong(KeyRangeIterator::getMaxKeys));
              int initialSize = rangeIterators.size();
              // all ranges will be included
              if (limit >= rangeIterators.size() || limit <= 0)
                  return buildIterator(statistics, rangeIterators);
  
              // Apply most selective iterators during intersection, because larger number of iterators will result lots of disk seek.
              Statistics selectiveStatistics = new IntersectionStatistics();
              isDisjoint = false;
              for (int i = rangeIterators.size() - 1; i >= 0 && i >= limit; i--)
                  FileUtils.closeQuietly(rangeIterators.remove(i));
  
              rangeIterators.forEach(range -> updateStatistics(selectiveStatistics, range));
  
              if (Tracing.isTracing())
                  Tracing.trace("Selecting {} {} of {} out of {} indexes",
                                rangeIterators.size(),
                                rangeIterators.size() > 1 ? "indexes with cardinalities" : "index with cardinality",
                                rangeIterators.stream().map(KeyRangeIterator::getMaxKeys).map(Object::toString).collect(Collectors.joining(", ")),
                                initialSize);
  
              return buildIterator(selectiveStatistics, rangeIterators);
          }
  
          public boolean isDisjoint()
          {
              return isDisjoint;
          }
  
          private KeyRangeIterator buildIterator(Statistics statistics, List<KeyRangeIterator> ranges)
          {
              // if the ranges are disjoint, or we have an intersection with an empty set,
              // we can simply return an empty iterator, because it's not going to produce any results.
              if (isDisjoint)
              {
                  FileUtils.closeQuietly(ranges);
                  onClose.run();
                  return KeyRangeIterator.empty();
              }
  
              if (ranges.size() == 1)
              {
                  KeyRangeIterator single = ranges.get(0);
                  single.setOnClose(onClose);
                  return single;
              }
  
              // Make sure intersection is supported on the ranges provided:
              PrimaryKey.Kind firstKind = null;
              
              for (KeyRangeIterator range : ranges)
              {
                  PrimaryKey key;
                  if(range.hasNext())
                      key = range.peek();
                  else
                      key = range.getMaximum();
  
                  if (key != null)
                      if (firstKind == null)
                          firstKind = key.kind();
                      else if (!firstKind.isIntersectable(key.kind()))
                          throw new IllegalArgumentException("Cannot intersect " + firstKind + " and " + key.kind() + " ranges!");
              }
  
              return new KeyRangeIntersectionIterator(statistics, ranges, onClose);
          }
  
          private void updateStatistics(Statistics statistics, KeyRangeIterator range)
          {
              statistics.update(range);
              isDisjoint |= isDisjointInternal(statistics.min, statistics.max, range);
          }
      }
  
      private static class IntersectionStatistics extends KeyRangeIterator.Builder.Statistics
      {
          private boolean empty = true;
  
          @Override
          public void update(KeyRangeIterator range)
          {
              // minimum of the intersection is the biggest minimum of individual iterators
              min = nullSafeMax(min, range.getMinimum());
              // maximum of the intersection is the smallest maximum of individual iterators
              max = nullSafeMin(max, range.getMaximum());
  
              // With STATIC keys, it is possible for the min to overtake the max, which must be corrected. 
              min = nullSafeMin(min, max);
  
              if (empty)
              {
                  empty = false;
                  count = range.getMaxKeys();
              }
              else
              {
                  count = Math.min(count, range.getMaxKeys());
              }
          }
      }
  
      @VisibleForTesting
      protected static boolean isDisjoint(KeyRangeIterator a, KeyRangeIterator b)
      {
          return isDisjointInternal(a.peek(), a.getMaximum(), b);
      }
  
      /**
       * Ranges are overlapping the following cases:
       * <p>
       *   * When they have a common subrange:
       * <p>
       *   min       b.current      max          b.max
       *   +---------|--------------+------------|
       * <p>
       *   b.current      min       max          b.max
       *   |--------------+---------+------------|
       * <p>
       *   min        b.current     b.max        max
       *   +----------|-------------|------------+
       * <p>
       *
       *  If either range is empty, they're disjoint.
       */
      private static boolean isDisjointInternal(PrimaryKey min, PrimaryKey max, KeyRangeIterator b)
      {
          return min == null || max == null || b.getMaxKeys() == 0
                 || min.compareTo(b.getMaximum(), false) > 0 || (b.hasNext() && b.peek().compareTo(max, false) > 0);
      }
  }
  