From 32030e4fa35813803ebaa93e1aa54964ae2b0cf5 Mon Sep 17 00:00:00 2001
From: Caleb Rackliffe <calebrackliffe@gmail.com>
Date: Tue, 29 Oct 2024 10:11:24 -0500
Subject: [PATCH] Reintroduce RestrictionSet#iterator() optimization around
 multi-column restrictions

patch by Caleb Rackliffe; reviewed by Benjamin Lerer for CASSANDRA-20034
---
 CHANGES.txt                                   |  1 +
 .../cql3/restrictions/RestrictionSet.java     | 62 +++++++++++++++++--
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index d81ffaebb0..fb216698c4 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,4 +1,5 @@
 5.1
+ * Reintroduce RestrictionSet#iterator() optimization around multi-column restrictions (CASSANDRA-20034)
  * Explicitly localize strings to Locale.US for internal implementation (CASSANDRA-19953)
  * Add -H option for human-friendly output in nodetool compactionhistory (CASSANDRA-20015)
  * Fix type check for referenced duration type for nested types (CASSANDRA-19890)
diff --git a/src/java/org/apache/cassandra/cql3/restrictions/RestrictionSet.java b/src/java/org/apache/cassandra/cql3/restrictions/RestrictionSet.java
index 9b3c760080..b14d3debef 100644
--- a/src/java/org/apache/cassandra/cql3/restrictions/RestrictionSet.java
+++ b/src/java/org/apache/cassandra/cql3/restrictions/RestrictionSet.java
@@ -23,12 +23,15 @@ import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashSet;
 import java.util.Iterator;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.NavigableMap;
 import java.util.Set;
 import java.util.TreeMap;
 
+import com.google.common.collect.AbstractIterator;
+import org.apache.commons.lang3.builder.ToStringBuilder;
+import org.apache.commons.lang3.builder.ToStringStyle;
+
 import org.apache.cassandra.cql3.QueryOptions;
 import org.apache.cassandra.cql3.functions.Function;
 import org.apache.cassandra.db.filter.RowFilter;
@@ -58,13 +61,18 @@ final class RestrictionSet implements Restrictions, Iterable<SingleRestriction>
     };
 
     private static final RestrictionSet EMPTY = new RestrictionSet(Collections.unmodifiableNavigableMap(new TreeMap<>(COLUMN_DEFINITION_COMPARATOR)),
-                                                                   false, false, false,false);
+                                                                   false, false, false, false,false);
 
     /**
      * The restrictions per column.
      */
     private final NavigableMap<ColumnMetadata, SingleRestriction> restrictions;
 
+    /**
+     * {@code true} if it contains multi-column restrictions, {@code false} otherwise.
+     */
+    private final boolean hasMultiColumnRestrictions;
+
     private final boolean hasSlice;
 
     private final boolean hasIn;
@@ -83,12 +91,14 @@ final class RestrictionSet implements Restrictions, Iterable<SingleRestriction>
     }
 
     private RestrictionSet(NavigableMap<ColumnMetadata, SingleRestriction> restrictions,
+                           boolean hasMultiColumnRestrictions,
                            boolean hasIn,
                            boolean hasSlice,
                            boolean hasAnn,
                            boolean needsFilteringOrIndexing)
     {
         this.restrictions = restrictions;
+        this.hasMultiColumnRestrictions = hasMultiColumnRestrictions;
         this.hasIn = hasIn;
         this.hasSlice = hasSlice;
         this.hasAnn = hasAnn;
@@ -184,6 +194,7 @@ final class RestrictionSet implements Restrictions, Iterable<SingleRestriction>
         boolean newNeedsFilteringOrIndexing = needsFilteringOrIndexing || restriction.needsFilteringOrIndexing();
 
         return new RestrictionSet(mergeRestrictions(newRestricitons, restriction),
+                                  hasMultiColumnRestrictions || restriction.isMultiColumn(),
                                   newHasIN,
                                   newHasSlice,
                                   newHasANN,
@@ -260,8 +271,8 @@ final class RestrictionSet implements Restrictions, Iterable<SingleRestriction>
     @Override
     public Iterator<SingleRestriction> iterator()
     {
-        // We need to eliminate duplicates in the case where we have multi-column restrictions.
-        return new LinkedHashSet<>(restrictions.values()).iterator();
+        Iterator<SingleRestriction> iterator = restrictions.values().iterator();
+        return hasMultiColumnRestrictions ? new DistinctIterator<>(iterator) : iterator;
     }
 
     @Override
@@ -301,4 +312,47 @@ final class RestrictionSet implements Restrictions, Iterable<SingleRestriction>
     {
         return restrictions.lastEntry().getValue();
     }
+
+    /**
+     * {@code Iterator} decorator that removes duplicates in an ordered one.
+     *
+     * @param <E> the iterator element type.
+     */
+    private static final class DistinctIterator<E> extends AbstractIterator<E>
+    {
+        /**
+         * The decorated iterator.
+         */
+        private final Iterator<E> iterator;
+
+        /**
+         * The previous element.
+         */
+        private E previous;
+
+        public DistinctIterator(Iterator<E> iterator)
+        {
+            this.iterator = iterator;
+        }
+
+        protected E computeNext()
+        {
+            while(iterator.hasNext())
+            {
+                E next = iterator.next();
+                if (!next.equals(previous))
+                {
+                    previous = next;
+                    return next;
+                }
+            }
+            return endOfData();
+        }
+    }
+
+    @Override
+    public String toString()
+    {
+        return ToStringBuilder.reflectionToString(this, ToStringStyle.SHORT_PREFIX_STYLE);
+    }
 }