Fix index issue with multiple keys in embedded documents

Ross Hansen · Ross Hansen · commit fe2adceb1fe4 · 2020-08-21T09:51:47.000-07:00
diff --git a/core/src/main/java/de/bwaldvogel/mongo/backend/CollectionUtils.java b/core/src/main/java/de/bwaldvogel/mongo/backend/CollectionUtils.java
@@ -37,22 +37,30 @@ public static <T> T getSingleElement(Iterable<T> iterable, Supplier<? extends Ru
         return value;
     }
 
-    static <T> List<List<T>> multiplyWithOtherElements(Collection<?> allValues, Collection<T> collectionValues) {
-        Assert.isTrue(allValues.contains(collectionValues), () -> "Expected " + collectionValues + " to be part of " + allValues);
-        List<List<T>> result = new ArrayList<>();
-        for (T collectionValue : collectionValues) {
-            List<T> values = new ArrayList<>();
+    static <T> List<List<T>> multiplyWithOtherElements(Collection<T> allValues, Collection<? extends Collection<T>> collectionValues) {
+        // Ex: multiplyWithOtherElements([abc, xyz, [L, M, S], [red, green, blue]], [[L, M, S], [red, green, blue]])
+        //         --> [[abc, xyz, L, red], [abc, xyz, M, green], [abc, xyz, S, blue]]
+        final int collectionValueSize = CollectionUtils.getElementAtPosition(collectionValues, 0).size();
+        for (Collection<T> collectionValue : collectionValues) {
+            Assert.isTrue(allValues.contains(collectionValue), () -> "Expected " + collectionValue + " to be part of " + allValues);
+            Assert.isTrue(collectionValue.size() == collectionValueSize, () -> "Expected " + collectionValue + " to be size " + collectionValueSize);
+        }
+
+        List<List<T>> result = new ArrayList<>(collectionValueSize);
+        for (int i = 0; i < collectionValueSize; i++) {
+            final ArrayList<T> newValues = new ArrayList();
 
             for (Object value : allValues) {
-                if (value == collectionValues) {
-                    values.add(collectionValue);
+                if (collectionValues.contains(value)) {
+                    newValues.add(CollectionUtils.getElementAtPosition((Collection<T>) value, i));
                 } else {
-                    values.add((T) value);
+                    newValues.add((T) value);
                 }
             }
 
-            result.add(values);
+            result.add(newValues);
         }
+
         return result;
     }
 
diff --git a/core/src/main/java/de/bwaldvogel/mongo/backend/Index.java b/core/src/main/java/de/bwaldvogel/mongo/backend/Index.java
@@ -65,20 +65,16 @@ Set<KeyValue> getKeyValues(Document document, boolean normalize) {
             valuesPerKey.replaceAll((key, value) -> Utils.normalizeValue(value));
         }
 
-        List<Collection<?>> collectionValues = valuesPerKey.values().stream()
+        List<Collection<Object>> collectionValues = valuesPerKey.values().stream()
             .filter(value -> value instanceof Collection)
-            .map(value -> (Collection<?>) value)
+            .map(value -> (Collection<Object>) value)
             .collect(Collectors.toList());
 
-        if (collectionValues.size() == 1) {
-            @SuppressWarnings("unchecked")
-            Collection<Object> collectionValue = (Collection<Object>) CollectionUtils.getSingleElement(collectionValues);
-            return CollectionUtils.multiplyWithOtherElements(valuesPerKey.values(), collectionValue).stream()
+        if (collectionValues.size() > 0) {
+            validateHasNoParallelArrays(document);
+            return CollectionUtils.multiplyWithOtherElements(valuesPerKey.values(), collectionValues).stream()
                 .map(KeyValue::new)
                 .collect(StreamUtils.toLinkedHashSet());
-        } else if (collectionValues.size() > 1) {
-            validateHasNoParallelArrays(document);
-            return collectCollectionValues(collectionValues);
         } else {
             return Collections.singleton(new KeyValue(valuesPerKey.values()));
         }
diff --git a/core/src/test/java/de/bwaldvogel/mongo/backend/CollectionUtilsTest.java b/core/src/test/java/de/bwaldvogel/mongo/backend/CollectionUtilsTest.java
@@ -64,17 +64,28 @@ void testMultiplyWithOtherElements() throws Exception {
         List<Object> values = Arrays.asList(1, 2);
         List<Object> collection = Arrays.asList("abc", "def", values);
 
-        assertThat(CollectionUtils.multiplyWithOtherElements(collection, values))
+        assertThat(CollectionUtils.multiplyWithOtherElements(collection, Arrays.asList(values)))
             .containsExactly(
                 Arrays.asList("abc", "def", 1),
                 Arrays.asList("abc", "def", 2)
             );
 
         assertThatExceptionOfType(IllegalArgumentException.class)
-            .isThrownBy(() -> CollectionUtils.multiplyWithOtherElements(collection, Arrays.asList(1, 2, 3)))
+            .isThrownBy(() -> CollectionUtils.multiplyWithOtherElements(collection, Arrays.asList(Arrays.asList(1, 2, 3))))
             .withMessage("Expected [1, 2, 3] to be part of [abc, def, [1, 2]]");
     }
 
+    @Test
+    void testMultiplyWithOtherElementsWrongSize() throws Exception {
+        List<Object> smallValues = Arrays.asList(1, 2);
+        List<Object> bigValues = Arrays.asList(1, 2, 3);
+        List<Object> collection = Arrays.asList("abc", "def", smallValues, bigValues);
+
+        assertThatExceptionOfType(IllegalArgumentException.class)
+            .isThrownBy(() -> CollectionUtils.multiplyWithOtherElements(collection, Arrays.asList(smallValues, bigValues)))
+            .withMessage("Expected [1, 2, 3] to be size 2");
+    }
+
     @Test
     void testGetElementAtPosition_list() throws Exception {
         List<String> list = Arrays.asList("a", "b", "c");
diff --git a/memory-backend/src/test/java/de/bwaldvogel/mongo/backend/memory/index/MemoryUniqueIndexTest.java b/memory-backend/src/test/java/de/bwaldvogel/mongo/backend/memory/index/MemoryUniqueIndexTest.java
@@ -72,6 +72,28 @@ void testGetKeyValues_multiKey_document_nested_objects() throws Exception {
         );
     }
 
+    @Test
+    void testGetKeyValues_multiKey_nested_objects_multiple_keys() throws Exception {
+        Index<?> index = new MemoryUniqueIndex("name", Arrays.asList(
+            new IndexKey("item", true),
+            new IndexKey("stock.size", true),
+            new IndexKey("stock.quantity", true)
+        ), false);
+
+        assertThat(index.getKeyValues(
+            new Document("stock", Arrays.asList(
+                new Document("size", "S").append("color", "red").append("quantity", 25),
+                new Document("size", "S").append("color", "blue").append("quantity", 10),
+                new Document("size", "M").append("color", "blue").append("quantity", 50)
+            ))
+                .append("item", "abc")
+        )).containsExactly(
+            new KeyValue("abc", "S", 25.0),
+            new KeyValue("abc", "S", 10.0),
+            new KeyValue("abc", "M", 50.0)
+        );
+    }
+
     private static Document jsonDocument(String json) {
         return convert(json(json));
     }
