Merge pull request #5 from bloxbean/multiget

feat: Multiget support
bloxbean · Dec 21, 2023 · 25eabf3 · 25eabf3
2 parents aba00b9 + e0248c5
commit 25eabf3
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 0 deletions.
diff --git a/core/src/main/java/com/bloxbean/rocks/types/collection/BaseDataType.java b/core/src/main/java/com/bloxbean/rocks/types/collection/BaseDataType.java
@@ -9,6 +9,7 @@
 import org.rocksdb.RocksIterator;
 import org.rocksdb.WriteBatch;
 
+import java.util.List;
 import java.util.Optional;
 
 /**
@@ -102,6 +103,19 @@ protected byte[] get(byte[] key) {
         }
     }
 
+    @SneakyThrows
+    protected List<byte[]> get(List<byte[]> keys) {
+        if (columnFamilyHandle != null) {
+            var columnFamilies = keys.stream()
+                    .map(key -> columnFamilyHandle)
+                    .toList();
+
+            return db.multiGetAsList(columnFamilies, keys);
+        } else {
+            return db.multiGetAsList(keys);
+        }
+    }
+
     @SneakyThrows
     protected void delete(byte[] key) {
         if (columnFamilyHandle != null) {

diff --git a/core/src/main/java/com/bloxbean/rocks/types/collection/RocksMap.java b/core/src/main/java/com/bloxbean/rocks/types/collection/RocksMap.java
@@ -6,6 +6,7 @@
 import lombok.NonNull;
 import org.rocksdb.WriteBatch;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -58,4 +59,8 @@ public Set<Map.Entry<K, V>> entries() {
     public ValueIterator<Map.Entry<K, V>> entriesIterator() {
         return super.entriesIterator(null);
     }
+
+    public List<V> multiGet(List<K> keys) {
+        return super.multiGet(null, keys);
+    }
 }
diff --git a/core/src/main/java/com/bloxbean/rocks/types/collection/RocksMultiMap.java b/core/src/main/java/com/bloxbean/rocks/types/collection/RocksMultiMap.java
@@ -59,6 +59,30 @@ public Optional<V> get(String ns, K key) {
         return Optional.of(valueSerializer.deserialize(valueBytes, valueType));
     }
 
+    public List<V> multiGet(String ns, List<K> keys) {
+        var metadata = getMetadata(ns);
+        if (metadata.isEmpty())
+            return Collections.emptyList();
+
+        var keysBytes = keys.stream()
+                .map(key -> getKey(metadata.get(), ns, key))
+                .toList();
+
+        List<byte[]> values = get(keysBytes);
+
+        if (values == null || values.size() == 0)
+            return Collections.emptyList();
+
+        return values.stream()
+                .map(value -> {
+                    if (value == null)
+                        return null;
+                    else
+                        return valueSerializer.deserialize(value, valueType);
+                })
+                .toList();
+    }
+
     public boolean contains(String ns, K key) {
         var metadata = getMetadata(ns);
         if (metadata.isEmpty())

diff --git a/core/src/test/java/com/bloxbean/rocks/types/collection/RocksMapTest.java b/core/src/test/java/com/bloxbean/rocks/types/collection/RocksMapTest.java
@@ -132,4 +132,16 @@ void entries_iterator() throws Exception {
                 Map.entry("key4", "value4")));
 
     }
+
+    @Test
+    void put_multiget() {
+        var map = new RocksMap<String, String>(rocksDBConfig, "testMap", String.class, String.class);
+        map.put("key1", "value1");
+        map.put("key2", "value2");
+        map.put("key3", "value3");
+        map.put("key4", "value4");
+
+        List<String> values = map.multiGet(List.of("key1", "key2", "key3", "key4"));
+        assertThat(values).isEqualTo(List.of("value1", "value2", "value3", "value4"));
+    }
 }
diff --git a/core/src/test/java/com/bloxbean/rocks/types/collection/RocksMultiMapTest.java b/core/src/test/java/com/bloxbean/rocks/types/collection/RocksMultiMapTest.java
@@ -149,4 +149,16 @@ void entries_iterator() throws Exception {
         assertThat(entries2).isEqualTo(List.of(Map.entry("key5", "value5"),
                 Map.entry("key6", "value6")));
     }
+
+    @Test
+    void put_multiget() {
+        var map = new RocksMultiMap<String, String>(rocksDBConfig, "testMap", String.class, String.class);
+        map.put("ns1", "key1", "value1");
+        map.put("ns1", "key2", "value2");
+        map.put("ns1", "key3", "value3");
+        map.put("ns1", "key4", "value4");
+
+        List<String> values = map.multiGet("ns1", List.of("key1", "key2", "key3", "key4"));
+        assertThat(values).isEqualTo(List.of("value1", "value2", "value3", "value4"));
+    }
 }