Fixes to the most-recent-unique functionality (#2392)

NationalSecurityAgency · May 30, 2024 · 3d27c1e · 3d27c1e
1 parent 78c14f4
commit 3d27c1e
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 13 deletions.
diff --git a/warehouse/query-core/src/main/java/datawave/query/planner/QueryOptionsSwitch.java b/warehouse/query-core/src/main/java/datawave/query/planner/QueryOptionsSwitch.java
@@ -66,6 +66,9 @@ public static void apply(Map<String,String> optionsMap, ShardQueryConfiguration
                     uniqueFields.setMostRecent(config.getUniqueFields().isMostRecent());
                     config.setUniqueFields(uniqueFields);
                     break;
+                case QueryParameters.MOST_RECENT_UNIQUE:
+                    config.getUniqueFields().setMostRecent(Boolean.parseBoolean(value));
+                    break;
                 case QueryParameters.EXCERPT_FIELDS:
                     ExcerptFields excerptFields = ExcerptFields.from(value);
                     config.setExcerptFields(excerptFields);

diff --git a/warehouse/query-core/src/main/java/datawave/query/tables/ShardQueryLogic.java b/warehouse/query-core/src/main/java/datawave/query/tables/ShardQueryLogic.java
@@ -913,15 +913,15 @@ protected void loadQueryParameters(ShardQueryConfiguration config, Query setting
             UniqueFields uniqueFields = UniqueFields.from(uniqueFieldsParam);
             // Only set the unique fields if we were actually given some
             if (!uniqueFields.isEmpty()) {
-                this.setUniqueFields(uniqueFields);
+                // preserve the most recent flag
+                uniqueFields.setMostRecent(config.getUniqueFields().isMostRecent());
                 config.setUniqueFields(uniqueFields);
             }
         }
 
         // Get the most recent flag
         String mostRecentUnique = settings.findParameter(QueryParameters.MOST_RECENT_UNIQUE).getParameterValue().trim();
         if (StringUtils.isNotBlank(mostRecentUnique)) {
-            this.getUniqueFields().setMostRecent(Boolean.valueOf(mostRecentUnique));
             config.getUniqueFields().setMostRecent(Boolean.valueOf(mostRecentUnique));
         }
 

diff --git a/...house/query-core/src/main/java/datawave/query/util/sortedset/RewritableSortedSetImpl.java b/...house/query-core/src/main/java/datawave/query/util/sortedset/RewritableSortedSetImpl.java
@@ -24,7 +24,7 @@
  */
 public class RewritableSortedSetImpl<E> implements RewritableSortedSet<E>, Cloneable {
     private static Logger log = Logger.getLogger(RewritableSortedSetImpl.class);
-    // using a map to enable replacement of the actual set member (see uses of collisionSelection)
+    // using a map to enable replacement of the actual set member (see uses of addResolvingCollisions)
     protected NavigableMap<E,E> set = null;
     // When the set contains X and we are adding Y where X == Y, then use this strategy
     // to decide which to keep.
@@ -128,18 +128,18 @@ public boolean contains(Object o) {
 
     @Override
     public Iterator<E> iterator() {
-        return set.keySet().iterator();
+        return set.values().iterator();
     }
 
     @Override
     public Object[] toArray() {
-        return set.keySet().toArray();
+        return set.values().toArray();
     }
 
     @SuppressWarnings({"unchecked"})
     @Override
     public <T> T[] toArray(T[] a) {
-        return set.keySet().toArray(a);
+        return set.values().toArray(a);
     }
 
     @Override
@@ -153,11 +153,16 @@ public E get(E e) {
     }
 
     private boolean addResolvingCollisions(E e) {
-        if ((rewriteStrategy != null) && set.containsKey(e) && rewriteStrategy.rewrite(set.get(e), e)) {
-            set.remove(e);
-        }
         // return true if this is a new element to the set
-        return (set.put(e, e) == null);
+        if (set.containsKey(e)) {
+            if ((rewriteStrategy != null) && rewriteStrategy.rewrite(set.get(e), e)) {
+                set.put(e, e);
+                return true;
+            }
+            return false;
+        }
+        set.put(e, e);
+        return true;
     }
 
     @Override
@@ -234,7 +239,7 @@ public E first() {
             QueryException qe = new QueryException(DatawaveErrorCode.FETCH_FIRST_ELEMENT_ERROR);
             throw (NoSuchElementException) (new NoSuchElementException().initCause(qe));
         } else {
-            return first;
+            return set.get(first);
         }
     }
 
@@ -250,13 +255,13 @@ public E last() {
             QueryException qe = new QueryException(DatawaveErrorCode.FETCH_LAST_ELEMENT_ERROR);
             throw (NoSuchElementException) (new NoSuchElementException().initCause(qe));
         } else {
-            return last;
+            return set.get(last);
         }
     }
 
     @Override
     public String toString() {
-        return set.toString();
+        return set.values().toString();
     }
 
     /**