perf: store a unique index as one id per key instead of a one-element list

A unique index kept the classic value -> [id] layout: a CopyOnWriteArrayList
per key that never held more than one element, allocated and copied on every
write, plus a size check standing in for the uniqueness test. The index now
stores the id itself (value -> id) in a map of its own, named with a "|unique"
suffix, and enforces uniqueness by comparing the stored id with the writer's:
another document under the key is a violation, the same document again is
not.

An index still in the list layout is migrated the first time it is accessed
and the legacy map is dropped, the way the composite layout migrates a
non-unique index. drop() removes whichever layouts exist without migrating
first. IndexMap exposes the single-id map to the scanner and the filters as
one-element lists, so the read path is unchanged; readSortKeys reads the
pairs directly. The map has its own name because the RocksDB adapter decodes
values by the declared type of the map they live in.

IndexManager.close(), clearAll() and dropIndexDescriptor() acted only on the
map name recorded in IndexMeta, which is the classic one. That already missed
the composite map of a non-unique index: after collection.clear() its rows
survived, and a query on that index returned the ids of the cleared
documents alongside the new ones (two live documents, four results). With
the unique layout it would have rejected the very keys the collection no
longer held. All three now cover every layout map an index can occupy.

Tests cover write, violation, same-document rewrite, remove by the right
and the wrong id, migration from the list layout, drop of both layouts,
sort keys, and clear() followed by fresh documents under the same keys on
both a non-unique and a unique index. The core and RocksDB suites pass.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: brettwooldridge

nitrite/src/main/java/org/dizitart/no2/collection/operation/IndexManager.java

@@ -28,7 +28,9 @@
 import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import static org.dizitart.no2.common.util.IndexUtils.deriveCompositeIndexMapName;
 import static org.dizitart.no2.common.util.IndexUtils.deriveIndexMapName;
+import static org.dizitart.no2.common.util.IndexUtils.deriveUniqueIndexMapName;
 import static org.dizitart.no2.common.util.IndexUtils.deriveIndexMetaMapName;
 
 /**
@@ -93,32 +95,57 @@ public void close() {
             Iterable<IndexMeta> indexMetas = indexMetaMap.values();
             for (IndexMeta indexMeta : indexMetas) {
                 if (indexMeta != null && indexMeta.getIndexDescriptor() != null) {
-                    String indexMapName = indexMeta.getIndexMap();
-                    NitriteMap<?, ?> indexMap = nitriteStore.openMap(indexMapName, Object.class, Object.class);
-                    indexMap.close();
+                    for (NitriteMap<?, ?> indexMap : existingLayoutMaps(indexMeta)) {
+                        indexMap.close();
+                    }
                 }
             }
-
             // close index meta
             indexMetaMap.close();
         }
     }
 
     public void clearAll() {
-        // close all index maps
+        // clear and close all index maps
         if (!indexMetaMap.isClosed() && !indexMetaMap.isDropped()) {
             Iterable<IndexMeta> indexMetas = indexMetaMap.values();
             for (IndexMeta indexMeta : indexMetas) {
                 if (indexMeta != null && indexMeta.getIndexDescriptor() != null) {
-                    String indexMapName = indexMeta.getIndexMap();
-                    NitriteMap<?, ?> indexMap = nitriteStore.openMap(indexMapName, Object.class, Object.class);
-                    indexMap.clear();
-                    indexMap.close();
+                    for (NitriteMap<?, ?> indexMap : existingLayoutMaps(indexMeta)) {
+                        indexMap.clear();
+                        indexMap.close();
+                    }
                 }
             }
         }
     }
 
+    /**
+     * The maps an index actually occupies in the store. {@link IndexMeta#getIndexMap()} records
+     * the classic map name, but a single-field index may instead live in the composite layout
+     * (non-unique) or the single-id layout (unique), each under a derived name of its own, and
+     * an index in mid-migration can briefly have two. Closing, clearing or dropping only the
+     * recorded map leaves the real one behind: after {@code clear()} its stale entries resolve
+     * to deleted documents, and a unique index rejects the very keys the collection no longer
+     * holds.
+     */
+    private List<NitriteMap<?, ?>> existingLayoutMaps(IndexMeta indexMeta) {
+        List<String> names = new ArrayList<>();
+        names.add(indexMeta.getIndexMap());
+        IndexDescriptor descriptor = indexMeta.getIndexDescriptor();
+        if (!descriptor.isCompoundIndex()) {
+            names.add(deriveCompositeIndexMapName(descriptor));
+            names.add(deriveUniqueIndexMapName(descriptor));
+        }
+        List<NitriteMap<?, ?>> maps = new ArrayList<>();
+        for (String name : names) {
+            if (nitriteStore.hasMap(name)) {
+                maps.add(nitriteStore.openMap(name, Object.class, Object.class));
+            }
+        }
+        return maps;
+    }
+
     /**
      * Is dirty index boolean.
      *
@@ -174,11 +201,10 @@ IndexDescriptor createIndexDescriptor(Fields fields, String indexType) {
     void dropIndexDescriptor(Fields fields) {
         IndexMeta meta = indexMetaMap.get(fields);
         if (meta != null && meta.getIndexDescriptor() != null) {
-            String indexMapName = meta.getIndexMap();
-            NitriteMap<?, ?> indexMap = nitriteStore.openMap(indexMapName, Object.class, Object.class);
-            indexMap.drop();
+            for (NitriteMap<?, ?> indexMap : existingLayoutMaps(meta)) {
+                indexMap.drop();
+            }
         }
-
         indexMetaMap.remove(fields);
         updateIndexDescriptorCache();
     }

nitrite/src/main/java/org/dizitart/no2/common/util/IndexUtils.java

@@ -48,6 +48,17 @@ public static String deriveCompositeIndexMapName(IndexDescriptor descriptor) {
         return deriveIndexMapName(descriptor) + INTERNAL_NAME_SEPARATOR + "composite";
     }
 
+    /**
+     * Derives the name of the map holding a unique index in its single-id layout, one
+     * {@code value -> id} entry per key.
+     *
+     * @param descriptor the index descriptor
+     * @return the map name
+     */
+    public static String deriveUniqueIndexMapName(IndexDescriptor descriptor) {
+        return deriveIndexMapName(descriptor) + INTERNAL_NAME_SEPARATOR + "unique";
+    }
+
     public static String deriveIndexMetaMapName(String collectionName) {
         return INDEX_META_PREFIX + INTERNAL_NAME_SEPARATOR + collectionName;
     }

nitrite/src/main/java/org/dizitart/no2/index/IndexMap.java

@@ -40,6 +40,8 @@ public class IndexMap {
     // (value, id) pairs (see IndexEntryKey). This IndexMap still presents the classic
     // value -> List<NitriteId> view to the scanner and the filters.
     private NitriteMap<IndexEntryKey, ?> compositeMap;
+    // single-id layout (unique index): values are NitriteIds, exposed as one-element lists
+    private boolean singleValued;
 
     @Getter
     @Setter
@@ -79,6 +81,24 @@ public static IndexMap composite(NitriteMap<IndexEntryKey, ?> compositeMap) {
         return new IndexMap(compositeMap, true);
     }
 
+    /**
+     * Instantiates an {@link IndexMap} over a unique index stored in the single-id layout
+     * ({@code value -> id}). The scanner and the filters expect a list of ids under every key,
+     * so each stored id is handed out as a one-element list.
+     *
+     * @param uniqueMap the backing map
+     * @return the index map
+     */
+    public static IndexMap unique(NitriteMap<DBValue, NitriteId> uniqueMap) {
+        IndexMap indexMap = new IndexMap(uniqueMap);
+        indexMap.singleValued = true;
+        return indexMap;
+    }
+
+    private static Object exposeValue(Object value, boolean singleValued) {
+        return singleValued && value instanceof NitriteId ? Collections.singletonList(value) : value;
+    }
+
     /**
      * Normalizes a key returned by the backing map to the {@link DBNull} singleton
      * when it represents the null key. Persistent stores deserialize the stored null
@@ -228,7 +248,7 @@ public Object get(DBValue dbValue) {
             return compositeGet(dbValue == null ? DBNull.getInstance() : dbValue);
         }
         if (nitriteMap != null) {
-            return nitriteMap.get(dbValue);
+            return exposeValue(nitriteMap.get(dbValue), singleValued);
         } else if (navigableMap != null) {
             return navigableMap.get(dbValue);
         }
@@ -284,10 +304,11 @@ public boolean hasNext() {
                 public Pair<DBValue, ?> next() {
                     Pair<DBValue, ?> next = entryIterator.next();
                     DBValue dbKey = next.getFirst();
+                    Object value = exposeValue(next.getSecond(), singleValued);
                     if (dbKey instanceof DBNull) {
-                        return new Pair<>(null, next.getSecond());
+                        return new Pair<>(null, value);
                     } else {
-                        return new Pair<>(dbKey, next.getSecond());
+                        return new Pair<>(dbKey, value);
                     }
                 }
             };

nitrite/src/main/java/org/dizitart/no2/index/NitriteIndex.java

@@ -132,8 +132,8 @@ default List<NitriteId> addNitriteIds(List<NitriteId> nitriteIds, FieldValues fi
             // ConcurrentModificationException. CopyOnWriteArrayList swaps its backing array
             // atomically on each mutation, so the background serializer always sees a stable
             // snapshot. Non-unique indexes avoid list values entirely via the composite layout
-            // (issue #1260); only unique indexes and the text index reach this path, where the
-            // per-key list is small enough that copy-on-write cost is negligible.
+            // (issue #1260) and unique indexes store their single id directly; only the text
+            // index still reaches this path.
             nitriteIds = new CopyOnWriteArrayList<>();
         }
 

nitrite/src/main/java/org/dizitart/no2/index/SingleFieldIndex.java

@@ -26,6 +26,7 @@
 import org.dizitart.no2.common.Fields;
 import org.dizitart.no2.common.tuples.Pair;
 import org.dizitart.no2.filters.ComparableFilter;
+import org.dizitart.no2.exceptions.UniqueConstraintException;
 import org.dizitart.no2.store.NitriteMap;
 import org.dizitart.no2.store.NitriteStore;
 
@@ -37,6 +38,7 @@
 
 import static org.dizitart.no2.common.util.IndexUtils.deriveCompositeIndexMapName;
 import static org.dizitart.no2.common.util.IndexUtils.deriveIndexMapName;
+import static org.dizitart.no2.common.util.IndexUtils.deriveUniqueIndexMapName;
 import static org.dizitart.no2.common.util.ObjectUtils.convertToObjectArray;
 
 /**
@@ -48,6 +50,7 @@ public class SingleFieldIndex implements NitriteIndex {
     private final IndexDescriptor indexDescriptor;
     private final NitriteStore<?> nitriteStore;
     private volatile boolean migrationChecked;
+    private volatile boolean uniqueMigrationChecked;
 
     /**
      * Instantiates a new {@link SingleFieldIndex}.
@@ -61,9 +64,9 @@ public SingleFieldIndex(IndexDescriptor indexDescriptor, NitriteStore<?> nitrite
     }
 
     /**
-     * The composite-key layout (issue #1260) is used for every non-unique index. Unique indexes
-     * keep the classic {@code value -> [id]} array layout because the uniqueness check relies on
-     * its single-array shape.
+     * The composite-key layout (issue #1260) is used for every non-unique index. A unique index
+     * has at most one id per key, so it stores that id directly ({@code value -> id}); the
+     * classic {@code value -> [id]} list layout it used before is migrated on first use.
      */
     private boolean useCompositeLayout() {
         return !isUnique();
@@ -78,10 +81,15 @@ public void write(FieldValues fieldValues) {
         Object element = fieldValues.get(firstField);
 
         if (!useCompositeLayout()) {
-            // unique indexes (and stores without comparable key ordering) keep the classic
-            // value -> [id] layout.
-            NitriteMap<DBValue, List<?>> indexMap = findIndexMap();
-            forEachElement(element, dbValue -> addIndexElement(indexMap, fieldValues, dbValue));
+            // one id per key: a violation is another document already holding the key
+            NitriteMap<DBValue, NitriteId> indexMap = findUniqueMap();
+            forEachElement(element, dbValue -> {
+                NitriteId existing = indexMap.get(dbValue);
+                if (existing != null && !existing.equals(fieldValues.getNitriteId())) {
+                    throw new UniqueConstraintException("Unique key constraint violation for " + fields);
+                }
+                indexMap.put(dbValue, fieldValues.getNitriteId());
+            });
         } else {
             // non-unique indexes use the composite-key layout: one O(log n) point write per
             // (value, id) pair, instead of an O(n) read-modify-write of a shared list (issue #1260)
@@ -100,8 +108,13 @@ public void remove(FieldValues fieldValues) {
         Object element = fieldValues.get(firstField);
 
         if (!useCompositeLayout()) {
-            NitriteMap<DBValue, List<?>> indexMap = findIndexMap();
-            forEachElement(element, dbValue -> removeIndexElement(indexMap, fieldValues, dbValue));
+            NitriteMap<DBValue, NitriteId> indexMap = findUniqueMap();
+            forEachElement(element, dbValue -> {
+                NitriteId existing = indexMap.get(dbValue);
+                if (existing != null && existing.equals(fieldValues.getNitriteId())) {
+                    indexMap.remove(dbValue);
+                }
+            });
         } else {
             NitriteMap<IndexEntryKey, Object> indexMap = findCompositeMap();
             forEachElement(element, dbValue ->
@@ -112,9 +125,10 @@ public void remove(FieldValues fieldValues) {
     @Override
     public void drop() {
         if (!useCompositeLayout()) {
-            NitriteMap<DBValue, List<?>> indexMap = findIndexMap();
-            indexMap.clear();
-            indexMap.drop();
+            // drop whichever layouts exist without migrating first; nothing being dropped
+            // needs converting
+            dropMapIfPresent(deriveUniqueIndexMapName(indexDescriptor), DBValue.class, NitriteId.class);
+            dropMapIfPresent(deriveIndexMapName(indexDescriptor), DBValue.class, ArrayList.class);
         } else {
             NitriteMap<IndexEntryKey, Object> indexMap = findCompositeMap();
             indexMap.clear();
@@ -129,7 +143,7 @@ public LinkedHashSet<NitriteId> findNitriteIds(FindPlan findPlan) {
 
         IndexMap iMap = useCompositeLayout()
             ? IndexMap.composite(findCompositeMap())
-            : new IndexMap(findIndexMap());
+            : IndexMap.unique(findUniqueMap());
         return scanIndex(findPlan, iMap);
     }
 
@@ -147,11 +161,9 @@ public List<Pair<DBValue, NitriteId>> readSortKeys(long collectionSize) {
                 keys.add(new Pair<>(key.getValue(), key.getNitriteId()));
             }
         } else {
-            for (Pair<DBValue, List<?>> entry : (Iterable<Pair<DBValue, List<?>>>) (Iterable<?>) findIndexMap().entries()) {
-                for (NitriteId nitriteId : (List<NitriteId>) entry.getSecond()) {
-                    if (!seen.add(nitriteId)) return null;
-                    keys.add(new Pair<>(entry.getFirst(), nitriteId));
-                }
+            for (Pair<DBValue, NitriteId> entry : findUniqueMap().entries()) {
+                if (!seen.add(entry.getSecond())) return null;
+                keys.add(new Pair<>(entry.getFirst(), entry.getSecond()));
             }
         }
 
@@ -180,25 +192,47 @@ private void forEachElement(Object element, java.util.function.Consumer<DBValue>
         }
     }
 
-    @SuppressWarnings("unchecked")
-    private void addIndexElement(NitriteMap<DBValue, List<?>> indexMap,
-                                 FieldValues fieldValues, DBValue element) {
-        List<NitriteId> nitriteIds = (List<NitriteId>) indexMap.get(element);
-        nitriteIds = addNitriteIds(nitriteIds, fieldValues);
-        indexMap.put(element, nitriteIds);
+    private NitriteMap<DBValue, NitriteId> findUniqueMap() {
+        migrateLegacyUniqueIndex();
+        return nitriteStore.openMap(deriveUniqueIndexMapName(indexDescriptor), DBValue.class, NitriteId.class);
     }
 
+    /**
+     * Rewrites a unique index left in the classic {@code value -> [id]} list layout into the
+     * single-id layout the first time the index is accessed, then drops the legacy map. The
+     * list layout paid a copy-on-write list per key for a list that never held more than one
+     * id. Idempotent and run once per index instance.
+     */
     @SuppressWarnings("unchecked")
-    private void removeIndexElement(NitriteMap<DBValue, List<?>> indexMap,
-                                    FieldValues fieldValues, DBValue element) {
-        List<NitriteId> nitriteIds = (List<NitriteId>) indexMap.get(element);
-        if (nitriteIds != null && !nitriteIds.isEmpty()) {
-            nitriteIds.remove(fieldValues.getNitriteId());
-            if (nitriteIds.size() == 0) {
-                indexMap.remove(element);
-            } else {
-                indexMap.put(element, nitriteIds);
+    private void migrateLegacyUniqueIndex() {
+        if (uniqueMigrationChecked) return;
+        synchronized (this) {
+            if (uniqueMigrationChecked) return;
+            String legacyName = deriveIndexMapName(indexDescriptor);
+            if (nitriteStore.hasMap(legacyName)) {
+                NitriteMap<DBValue, List<?>> legacy = findIndexMap();
+                if (!legacy.isEmpty()) {
+                    NitriteMap<DBValue, NitriteId> unique = nitriteStore.openMap(
+                        deriveUniqueIndexMapName(indexDescriptor), DBValue.class, NitriteId.class);
+                    for (Pair<DBValue, List<?>> entry : (Iterable<Pair<DBValue, List<?>>>) (Iterable<?>) legacy.entries()) {
+                        List<NitriteId> nitriteIds = (List<NitriteId>) entry.getSecond();
+                        if (nitriteIds != null && !nitriteIds.isEmpty()) {
+                            unique.put(entry.getFirst(), nitriteIds.get(0));
+                        }
+                    }
+                }
+                legacy.clear();
+                legacy.drop();
             }
+            uniqueMigrationChecked = true;
+        }
+    }
+
+    private void dropMapIfPresent(String mapName, Class<?> keyType, Class<?> valueType) {
+        if (nitriteStore.hasMap(mapName)) {
+            NitriteMap<?, ?> map = nitriteStore.openMap(mapName, keyType, valueType);
+            map.clear();
+            map.drop();
         }
     }