optimize: extract the common logic

Author: lokidundun

hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBTable.java

@@ -210,6 +210,18 @@ protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
             return this.queryById(session, ids.iterator().next());
         }
 
+        // NOTE: this will lead to lazy create rocksdb iterator
+        return BackendColumnIterator.wrap(new FlatMapperIterator<>(
+                ids.iterator(), id -> this.queryById(session, id)
+        ));
+    }
+
+    protected BackendColumnIterator queryByIdsWithGet(RocksDBSessions.Session session,
+                                                      Collection<Id> ids) {
+        if (ids.size() == 1) {
+            return this.queryById(session, ids.iterator().next());
+        }
+
         if (!session.hasChanges()) {
             return this.getByIds(session, ids);
         }

hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBTables.java

@@ -19,6 +19,7 @@
 
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.util.Collection;
 import java.util.List;
 
 import org.apache.hugegraph.backend.id.Id;
@@ -177,6 +178,12 @@ public Vertex(String database) {
         protected BackendColumnIterator queryById(RocksDBSessions.Session session, Id id) {
             return this.getById(session, id);
         }
+
+        @Override
+        protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
+                                                   Collection<Id> ids) {
+            return this.queryByIdsWithGet(session, ids);
+        }
     }
 
     public static class Edge extends RocksDBTable {
@@ -200,6 +207,12 @@ public static Edge in(String database) {
         protected BackendColumnIterator queryById(RocksDBSessions.Session session, Id id) {
             return this.getById(session, id);
         }
+
+        @Override
+        protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
+                                                   Collection<Id> ids) {
+            return this.queryByIdsWithGet(session, ids);
+        }
     }
 
     public static class IndexTable extends RocksDBTable {

hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/rocksdb/RocksDBTableQueryByIdsTest.java

@@ -171,19 +171,27 @@ public boolean hasChanges() {
                 return true;
             }
 
+            @Override
+            public byte[] get(String table, byte[] key) {
+                throw new AssertionError(
+                        "reads should not be performed when hasChanges");
+            }
+
             @Override
             public BackendColumnIterator get(String table, List<byte[]> keys) {
                 throw new AssertionError(
                         "multi-get should not be called when hasChanges");
             }
         };
 
-        BackendColumnIterator iter = this.vertexTable.queryByIds(mockSession, ids);
-
-        Map<String, String> results = toResultMap(iter);
-        Assert.assertEquals(2, results.size());
-        Assert.assertEquals("value1", results.get("v1"));
-        Assert.assertEquals("value2", results.get("v2"));
+        try {
+            BackendColumnIterator iter = this.vertexTable.queryByIds(mockSession, ids);
+            // FlatMapperIterator is lazy; trigger evaluation to hit the mock
+            iter.hasNext();
+            Assert.fail("queryByIds should fail when session has pending changes");
+        } catch (AssertionError e) {
+            Assert.assertTrue(e.getMessage().contains("hasChanges"));
+        }
     }
 
     private Map<String, String> toResultMap(BackendColumnIterator iter) {