InMemoryDataShard: Allow to replace TreeMap with an explicitly sorted LinkedHashMap to reduce comparison count

`TreeMap` requires log(N) comparisons for every `Map.get()` and every `Map.put()` operation.
Due to the implementation of in-memory storage, this **not** cheap.

If you have a massive amount of `insert()`s/`save()`s and only occasionally get the list of sorted entities,
you will likely gain massive performance benefits from using an alternative in-memory data shard map
implementation using `LinkedHashMap` (even if sorted on `insert()`, but especially if sorted on reads).

To switch map implementation, set the `tech.ydb.yoj.repository.test.inmemory.impl` system property
to `oninsert` (for sorting data on `insert()`) or `onget` (for sorting data on first read.)

---------

Co-authored-by: Sergei Galiamichev <gsv@nebius.com>

Author: g-sg-v

repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/EntityIdMap.java

@@ -0,0 +1,90 @@
+package tech.ydb.yoj.repository.test.inmemory;
+
+import tech.ydb.yoj.repository.db.Entity;
+
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * TreeMap<T, E> - is very expensive due to very complicated comparison log(n)
+ */
+class EntityIdMap<T extends Entity<T>, E> extends AbstractMap<Entity.Id<T>, E> {
+    private final Comparator<Entity.Id<T>> comparator;
+    private Map<Entity.Id<T>, E> entries = new LinkedHashMap<>(); // probably not the best here but very simple
+
+    public EntityIdMap(Comparator<Entity.Id<T>> comparator) {
+        this.comparator = comparator;
+    }
+
+    public EntityIdMap(EntityIdMap<T, E> entityLines) {
+        this(entityLines.comparator);
+        entries = new LinkedHashMap<>(entityLines);
+    }
+
+    @Override
+    public Set<Entry<Entity.Id<T>, E>> entrySet() {
+        return entries.entrySet();
+    }
+
+    @Override
+    public E get(Object key) {
+        return entries.get(key);
+    }
+
+    @Override
+    public E put(Entity.Id<T> key, E value) {
+        List<Entity.Id<T>> temp = new ArrayList<>(entries.keySet());
+        int index = Collections.binarySearch(temp, key, comparator);
+        Map<Entity.Id<T>, E> result = new LinkedHashMap<>();
+        E oldValue = null;
+        int i = 0;
+        for (Entry<Entity.Id<T>, E> entry : entries.entrySet()) {
+            if (index < 0 && -index - 1 == i) {
+                result.put(key, value);
+                result.put(entry.getKey(), entry.getValue());
+            } else if (index == i) {
+                oldValue = entry.getValue();
+                result.put(key, value);
+            } else {
+                result.put(entry.getKey(), entry.getValue());
+            }
+            i++;
+        }
+        if (-index - 1 == result.size()) {
+            result.put(key, value); // add last (item i.e first item or item with max value to existing list)
+        }
+        entries = result;
+        return oldValue;
+    }
+
+    @Override
+    public E remove(Object key) {
+        return entries.remove(key);
+    }
+
+    @Override
+    public void clear() {
+        entries.clear();
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+        return entries.containsKey(key);
+    }
+
+    @Override
+    public boolean containsValue(Object value) {
+        return entries.containsValue(value);
+    }
+
+    @Override
+    public int size() {
+        return entries.size();
+    }
+}

repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/EntityIdMapOnGet.java

@@ -0,0 +1,74 @@
+package tech.ydb.yoj.repository.test.inmemory;
+
+import tech.ydb.yoj.repository.db.Entity;
+
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+class EntityIdMapOnGet<T extends Entity<T>, E> extends AbstractMap<Entity.Id<T>, E> {
+    private final Comparator<Entity.Id<T>> comparator;
+    private final Map<Entity.Id<T>, E> entries = new LinkedHashMap<>();
+    private Set<Entry<Entity.Id<T>, E>> cachedEntrySet = null;
+
+    public EntityIdMapOnGet(Comparator<Entity.Id<T>> comparator) {
+        this.comparator = comparator;
+    }
+
+    @Override
+    public Set<Entry<Entity.Id<T>, E>> entrySet() {
+        if (cachedEntrySet == null) {
+            List<Entry<Entity.Id<T>, E>> sorted = new ArrayList<>(entries.entrySet());
+            sorted.sort(Comparator.comparing(Entry::getKey, comparator));
+            cachedEntrySet = new LinkedHashSet<>(sorted);
+        }
+        return cachedEntrySet;
+    }
+
+    @Override
+    public E get(Object key) {
+        return entries.get(key);
+    }
+
+    @Override
+    public E put(Entity.Id<T> key, E value) {
+        E old = entries.put(key, value);
+        cachedEntrySet = null; // invalidate cache
+        return old;
+    }
+
+    @Override
+    public E remove(Object key) {
+        E removed = entries.remove(key);
+        if (removed != null) {
+            cachedEntrySet = null; // invalidate cache
+        }
+        return removed;
+    }
+
+    @Override
+    public void clear() {
+        entries.clear();
+        cachedEntrySet = null;
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+        return entries.containsKey(key);
+    }
+
+    @Override
+    public boolean containsValue(Object value) {
+        return entries.containsValue(value);
+    }
+
+    @Override
+    public int size() {
+        return entries.size();
+    }
+}

repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryDataShard.java

@@ -21,15 +21,18 @@
 import java.util.TreeMap;
 
 /*package*/ final class InMemoryDataShard<T extends Entity<T>> {
+    private static final String DEFAULT_MAP_IMPLEMENTATION = "treemap";
+    private static final String MAP_IMPLEMENTATION = System.getProperty("tech.ydb.yoj.repository.test.inmemory.impl", DEFAULT_MAP_IMPLEMENTATION);
+
     private final TableDescriptor<T> tableDescriptor;
     private final EntitySchema<T> schema;
-    private final TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines;
+    private final Map<Entity.Id<T>, InMemoryEntityLine> entityLines;
     private final Map<Long, Set<Entity.Id<T>>> uncommited = new HashMap<>();
 
     private InMemoryDataShard(
             TableDescriptor<T> tableDescriptor,
             EntitySchema<T> schema,
-            TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines
+            Map<Entity.Id<T>, InMemoryEntityLine> entityLines
     ) {
         this.tableDescriptor = tableDescriptor;
         this.schema = schema;
@@ -44,12 +47,17 @@ public InMemoryDataShard(TableDescriptor<T> tableDescriptor) {
         );
     }
 
-    private static <T extends Entity<T>> TreeMap<Entity.Id<T>, InMemoryEntityLine> createEmptyLines(Class<T> type) {
+    private static <T extends Entity<T>> Map<Entity.Id<T>, InMemoryEntityLine> createEmptyLines(Class<T> type) {
+        if ("oninsert".equals(MAP_IMPLEMENTATION)) {
+            return new EntityIdMap<>(EntityIdSchema.getIdComparator(type));
+        } else if ("onget".equals(MAP_IMPLEMENTATION)) {
+            return new EntityIdMapOnGet<>(EntityIdSchema.getIdComparator(type));
+        }
         return new TreeMap<>(EntityIdSchema.getIdComparator(type));
     }
 
     public synchronized InMemoryDataShard<T> createSnapshot() {
-        TreeMap<Entity.Id<T>, InMemoryEntityLine> snapshotLines = createEmptyLines(tableDescriptor.entityType());
+        Map<Entity.Id<T>, InMemoryEntityLine> snapshotLines = createEmptyLines(tableDescriptor.entityType());
         for (Map.Entry<Entity.Id<T>, InMemoryEntityLine> entry : entityLines.entrySet()) {
             snapshotLines.put(entry.getKey(), entry.getValue().createSnapshot());
         }