noticket: Get rid of most calls to Entity.Id.getType() and use better and simpler logic for "default sort order"

...because `Entity.Id.getType()` uses `TypeToken` which is **extremely slow** in tight loops
when run via Bazel JUnit5 tests, which have a non-null `SecurityManager` installed
(`SecurityManager` is to prevent `System.exit()` calls from tests, but it happens to check
every reflective access made by `TypeToken`, of which there are **plenty**.)

Author: nvamelichev

databind/src/main/java/tech/ydb/yoj/InternalApi.java

@@ -18,7 +18,7 @@
  * Annotates internal YOJ implementation details (classes, interfaces, methods, fields, constants, etc.) that need to be
  * {@code public} to be used from different and/or multiple packages, but are not stable even across minor YOJ releases.
  * <p>Non-{@code public} (e.g., package-private) classes, interfaces, methods, fields, constants etc. in YOJ are assumed
- * to be internal implementation details regardless of the presence of an {@code &#64;InternalApi} annotation on them.
+ * to be internal implementation details regardless of the presence of an {@code @InternalApi} annotation on them.
  */
 @Target({TYPE, FIELD, METHOD, PARAMETER, CONSTRUCTOR, ANNOTATION_TYPE, PACKAGE, MODULE, RECORD_COMPONENT})
 @Retention(SOURCE)

databind/src/main/java/tech/ydb/yoj/databind/schema/reflect/StdReflector.java

@@ -49,7 +49,7 @@ public ReflectType<?> reflectFieldType(Type genericType, FieldValueType bindingT
     }
 
     private ReflectType<?> reflectFor(Type type, FieldValueType fvt) {
-        Class<?> rawType = TypeToken.of(type).getRawType();
+        Class<?> rawType = type instanceof Class<?> clazz ? clazz : TypeToken.of(type).getRawType();
         for (TypeFactory m : matchers) {
             if (m.matches(rawType, fvt)) {
                 return m.create(this, rawType, fvt);

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

@@ -2,7 +2,6 @@
 
 import tech.ydb.yoj.databind.schema.Schema;
 import tech.ydb.yoj.repository.db.Entity;
-import tech.ydb.yoj.repository.db.EntityIdSchema;
 import tech.ydb.yoj.repository.db.EntitySchema;
 import tech.ydb.yoj.repository.db.Range;
 import tech.ydb.yoj.repository.db.Table;
@@ -18,18 +17,19 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.SortedMap;
 import java.util.TreeMap;
 
 /*package*/ final class InMemoryDataShard<T extends Entity<T>> {
     private final TableDescriptor<T> tableDescriptor;
     private final EntitySchema<T> schema;
-    private final TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines;
+    private final SortedMap<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
+            SortedMap<Entity.Id<T>, InMemoryEntityLine> entityLines
     ) {
         this.tableDescriptor = tableDescriptor;
         this.schema = schema;
@@ -39,18 +39,24 @@ private InMemoryDataShard(
     public InMemoryDataShard(TableDescriptor<T> tableDescriptor) {
         this(
                 tableDescriptor,
-                EntitySchema.of(tableDescriptor.entityType()),
-                createEmptyLines(tableDescriptor.entityType())
+                EntitySchema.of(tableDescriptor.entityType())
         );
     }
 
-    private static <T extends Entity<T>> TreeMap<Entity.Id<T>, InMemoryEntityLine> createEmptyLines(Class<T> type) {
-        return new TreeMap<>(EntityIdSchema.getIdComparator(type));
+    private InMemoryDataShard(
+            TableDescriptor<T> tableDescriptor,
+            EntitySchema<T> schema
+    ) {
+        this(tableDescriptor, schema, createEmptyLines(schema));
+    }
+
+    private static <T extends Entity<T>> SortedMap<Entity.Id<T>, InMemoryEntityLine> createEmptyLines(EntitySchema<T> schema) {
+        return new TreeMap<>(schema.getIdSchema());
     }
 
     public synchronized InMemoryDataShard<T> createSnapshot() {
-        TreeMap<Entity.Id<T>, InMemoryEntityLine> snapshotLines = createEmptyLines(tableDescriptor.entityType());
-        for (Map.Entry<Entity.Id<T>, InMemoryEntityLine> entry : entityLines.entrySet()) {
+        var snapshotLines = createEmptyLines(schema);
+        for (var entry : entityLines.entrySet()) {
             snapshotLines.put(entry.getKey(), entry.getValue().createSnapshot());
         }
         return new InMemoryDataShard<>(tableDescriptor, schema, snapshotLines);

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

@@ -3,6 +3,7 @@
 import com.google.common.base.Stopwatch;
 import com.google.common.collect.Iterables;
 import lombok.Getter;
+import tech.ydb.yoj.DeprecationWarnings;
 import tech.ydb.yoj.repository.BaseDb;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.RepositoryTransaction;
@@ -57,16 +58,21 @@ private long getVersion() {
 
     @Override
     public <T extends Entity<T>> Table<T> table(Class<T> c) {
-        return new InMemoryTable<>(getMemory(c));
+        return new InMemoryTable<>(this, c);
     }
 
     @Override
     public <T extends Entity<T>> Table<T> table(TableDescriptor<T> tableDescriptor) {
         return new InMemoryTable<>(this, tableDescriptor);
     }
 
-    @Deprecated // use other constructor of InMemoryTable
+    /**
+     * @deprecated {@code DbMemory} and this method will be removed in YOJ 3.0.0.
+     */
+    @Deprecated(forRemoval = true)
     public final <T extends Entity<T>> InMemoryTable.DbMemory<T> getMemory(Class<T> c) {
+        DeprecationWarnings.warnOnce("InMemoryTable.getMemory",
+                "InMemoryTable.getMemory(Class<T>) will be removed in YOJ 3.0.0. Please stop using this method"); 
         return new InMemoryTable.DbMemory<>(c, this);
     }
 

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

@@ -3,6 +3,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
+import tech.ydb.yoj.DeprecationWarnings;
 import tech.ydb.yoj.databind.expression.FilterExpression;
 import tech.ydb.yoj.databind.expression.OrderExpression;
 import tech.ydb.yoj.databind.schema.ObjectSchema;
@@ -14,6 +15,7 @@
 import tech.ydb.yoj.repository.db.Range;
 import tech.ydb.yoj.repository.db.Table;
 import tech.ydb.yoj.repository.db.TableDescriptor;
+import tech.ydb.yoj.repository.db.TableQueryBuilder;
 import tech.ydb.yoj.repository.db.TableQueryImpl;
 import tech.ydb.yoj.repository.db.ViewSchema;
 import tech.ydb.yoj.repository.db.cache.FirstLevelCache;
@@ -32,19 +34,28 @@
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toUnmodifiableMap;
 import static java.util.stream.Collectors.toUnmodifiableSet;
+import static tech.ydb.yoj.repository.db.TableQueryImpl.getEntityByIdComparator;
 
 public class InMemoryTable<T extends Entity<T>> implements Table<T> {
     private final EntitySchema<T> schema;
     private final TableDescriptor<T> tableDescriptor;
     private final InMemoryRepositoryTransaction transaction;
 
-    @Deprecated // Don't use DbMemory, use other constructor instead
+    /**
+     * @deprecated {@code DbMemory} and this constructor will be removed in YOJ 3.0.0.
+     * Please use other constructors instead.
+     */
+    @Deprecated(forRemoval = true)
     public InMemoryTable(DbMemory<T> memory) {
         this(memory.transaction(), memory.type());
+        DeprecationWarnings.warnOnce("new InMemoryTable(DbMemory)",
+                "Please do not use the InMemoryTable(DbMemory<T>) constructor, it will be removed in YOJ 3.0.0");
     }
 
     public InMemoryTable(InMemoryRepositoryTransaction transaction, Class<T> type) {
-        this(transaction, TableDescriptor.from(EntitySchema.of(type)));
+        this.schema = EntitySchema.of(type);
+        this.tableDescriptor = TableDescriptor.from(schema);
+        this.transaction = transaction;
     }
 
     public InMemoryTable(InMemoryRepositoryTransaction transaction, TableDescriptor<T> tableDescriptor) {
@@ -53,6 +64,11 @@ public InMemoryTable(InMemoryRepositoryTransaction transaction, TableDescriptor<
         this.transaction = transaction;
     }
 
+    @Override
+    public TableQueryBuilder<T> query() {
+        return new TableQueryBuilder<>(this, schema);
+    }
+    
     @Override
     public List<T> findAll() {
         transaction.getWatcher().markTableRead(tableDescriptor);
@@ -103,7 +119,7 @@ public List<T> find(
             @Nullable Long offset
     ) {
         // NOTE: InMemoryTable doesn't handle index.
-        return InMemoryQueries.find(() -> findAll().stream(), filter, orderBy, limit, offset);
+        return TableQueryImpl.find(() -> findAll().stream(), schema, filter, orderBy, limit, offset);
     }
 
     @Override
@@ -173,8 +189,8 @@ public TableDescriptor<T> getTableDescriptor() {
 
     @Override
     public T find(Entity.Id<T> id) {
-        if (id.isPartial()) {
-            throw new IllegalArgumentException("Cannot use partial id in find method");
+        if (TableQueryImpl.isPartialId(id, schema)) {
+            throw new IllegalArgumentException("Cannot use partial ID in Table.find() method");
         }
         return transaction.getTransactionLocal().firstLevelCache(tableDescriptor).get(id, __ -> {
             markKeyRead(id);
@@ -185,13 +201,13 @@ public T find(Entity.Id<T> id) {
 
     @Override
     public <ID extends Entity.Id<T>> List<T> find(Set<ID> ids) {
-        return TableQueryImpl.find(this, getFirstLevelCache(), ids);
+        return TableQueryImpl.find(this, schema, getFirstLevelCache(), ids);
     }
 
     @Override
     public <V extends View> V find(Class<V> viewType, Entity.Id<T> id) {
-        if (id.isPartial()) {
-            throw new IllegalArgumentException("Cannot use partial id in find method");
+        if (TableQueryImpl.isPartialId(id, schema)) {
+            throw new IllegalArgumentException("Cannot use partial ID in Table.find() method");
         }
 
         FirstLevelCache<T> cache = transaction.getTransactionLocal().firstLevelCache(tableDescriptor);
@@ -208,10 +224,13 @@ public <V extends View> V find(Class<V> viewType, Entity.Id<T> id) {
     @Override
     @SuppressWarnings("unchecked")
     public <ID extends Entity.Id<T>> List<T> find(Range<ID> range) {
-        transaction.getWatcher().markRangeRead(tableDescriptor, range);
+        Preconditions.checkArgument(range.getType() == schema.getIdSchema(),
+                "ID schema mismatch: Range was constructed with a different ID schema than the YdbTable");
+
+        markRangeRead(range);
         return findAll0().stream()
                 .filter(e -> range.contains((ID) e.getId()))
-                .sorted(EntityIdSchema.SORT_ENTITY_BY_ID)
+                .sorted(getEntityByIdComparator(schema))
                 .collect(toList());
     }
 
@@ -232,7 +251,7 @@ public <V extends View, ID extends Entity.Id<T>> List<V> find(Class<V> viewType,
 
     @Override
     public <V extends View, ID extends Entity.Id<T>> List<V> find(Class<V> viewType, Set<ID> ids) {
-        return find(viewType, ids, null, EntityExpressions.defaultOrder(getType()), null);
+        return find(viewType, ids, null, EntityExpressions.defaultOrder(schema), null);
     }
 
     @Override
@@ -404,9 +423,10 @@ private boolean isPrefixedFields(List<String> keyFields, Set<String> fields) {
 
     private <ID extends Entity.Id<T>> void markKeyRead(ID id) {
         EntityIdSchema<Entity.Id<T>> idSchema = schema.getIdSchema();
-        if (idSchema.flattenFieldNames().size() != idSchema.flatten(id).size()) {
+        Map<String, Object> eqMap = idSchema.flatten(id);
+        if (idSchema.flattenFieldNames().size() != eqMap.size()) {
             // Partial key, will throw error when not searching by PK prefix
-            transaction.getWatcher().markRangeRead(tableDescriptor, Range.create(id, id));
+            transaction.getWatcher().markRangeRead(tableDescriptor, Range.create(idSchema, eqMap));
         } else {
             transaction.getWatcher().markRowRead(tableDescriptor, id);
         }
@@ -475,7 +495,7 @@ public <ID extends Entity.Id<T>> Stream<T> streamPartial(ID partial, int batchSi
         Preconditions.checkArgument(1 <= batchSize && batchSize <= 5000,
                 "batchSize must be in range [1, 5000], got %s", batchSize);
 
-        Range<ID> range = partial == null ? null : Range.create(partial);
+        Range<ID> range = rangeForPartialId(partial);
         markRangeRead(range);
 
         return streamPartial0(range);
@@ -508,20 +528,26 @@ public <ID extends Entity.Id<T>> Stream<ID> streamPartialIds(ID partial, int bat
         Preconditions.checkArgument(1 <= batchSize && batchSize <= 10000,
                 "batchSize must be in range [1, 10000], got %s", batchSize);
 
-        Range<ID> range = partial == null ? null : Range.create(partial);
+        Range<ID> range = rangeForPartialId(partial);
         markRangeRead(range);
 
         return streamPartial0(range).map(e -> (ID) e.getId());
     }
 
-    private <ID extends Entity.Id<T>> void markRangeRead(Range<ID> range) {
+    private <ID extends Entity.Id<T>> void markRangeRead(@Nullable Range<ID> range) {
         if (range == null) {
             transaction.getWatcher().markTableRead(tableDescriptor);
         } else {
             transaction.getWatcher().markRangeRead(tableDescriptor, range);
         }
     }
 
+    @Nullable
+    private <ID extends Entity.Id<T>> Range<ID> rangeForPartialId(ID partial) {
+        EntityIdSchema<ID> idSchema = schema.getIdSchema();
+        return partial == null ? null : Range.create(idSchema, idSchema.flatten(partial));
+    }
+
     private <ID extends Entity.Id<T>> Stream<T> readTableStream(ReadTableParams<ID> params) {
         if (!transaction.getOptions().getIsolationLevel().isReadOnly()) {
             throw new IllegalTransactionIsolationLevelException("readTable", transaction.getOptions().getIsolationLevel());
@@ -533,7 +559,7 @@ private <ID extends Entity.Id<T>> Stream<T> readTableStream(ReadTableParams<ID>
                 .stream()
                 .filter(e -> readTableFilter(e, params));
         if (params.isOrdered()) {
-            stream = stream.sorted(EntityIdSchema.SORT_ENTITY_BY_ID);
+            stream = stream.sorted(getEntityByIdComparator(schema));
         }
         if (params.getRowLimit() > 0) {
             stream = stream.limit(params.getRowLimit());
@@ -542,18 +568,20 @@ private <ID extends Entity.Id<T>> Stream<T> readTableStream(ReadTableParams<ID>
     }
 
     private <ID extends Entity.Id<T>> boolean readTableFilter(T e, ReadTableParams<ID> params) {
+        EntityIdSchema<ID> idSchema = schema.getIdSchema();
+
         @SuppressWarnings("unchecked")
         ID id = (ID) e.getId();
         ID from = params.getFromKey();
         if (from != null) {
-            int compare = EntityIdSchema.ofEntity(id.getType()).compare(id, from);
+            int compare = idSchema.compare(id, from);
             if (params.isFromInclusive() ? compare < 0 : compare <= 0) {
                 return false;
             }
         }
         ID to = params.getToKey();
         if (to != null) {
-            int compare = EntityIdSchema.ofEntity(id.getType()).compare(id, to);
+            int compare = idSchema.compare(id, to);
             return params.isToInclusive() ? compare <= 0 : compare < 0;
         }
         return true;
@@ -587,7 +615,12 @@ private static <V extends Table.View, T extends Entity<T>> V toView(
     }
 
 
-    @Deprecated // Legacy. Using only for creating InMemoryTable. Use constructor of InMemoryTable instead
+    /**
+     * @deprecated Legacy class, used only for creating {@code InMemoryTable}. 
+     * This class will be removed in YOJ 3.0.0.
+     * Please use other constructors of {@code InMemoryTable} instead.
+     */
+    @Deprecated(forRemoval = true)
     public record DbMemory<T extends Entity<T>>(
             Class<T> type,
             InMemoryRepositoryTransaction transaction

repository-inmemory/src/test/java/tech/ydb/yoj/repository/test/inmemory/TestInMemoryRepository.java

@@ -105,7 +105,7 @@ public SupabubbleTable supabubbles() {
 
         @Override
         public Supabubble2Table supabubbles2() {
-            return new Supabubble2InMemoryTable(getMemory(Supabubble2.class));
+            return new Supabubble2InMemoryTable(this, Supabubble2.class);
         }
 
         @Override
@@ -145,8 +145,8 @@ public Table<MultiWrappedEntity2> multiWrappedEntities2() {
     }
 
     private static class Supabubble2InMemoryTable extends InMemoryTable<Supabubble2> implements TestEntityOperations.Supabubble2Table {
-        public Supabubble2InMemoryTable(DbMemory<Supabubble2> memory) {
-            super(memory);
+        public Supabubble2InMemoryTable(InMemoryRepositoryTransaction transaction, Class<Supabubble2> type) {
+            super(transaction, type);
         }
     }