Pull up all matches of a value through value that has multiple matches

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/combinatorics/PartiallyOrderedSet.java

@@ -29,11 +29,13 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Multimap;
 import com.google.common.collect.SetMultimap;
 import com.google.common.collect.Sets;
 
 import javax.annotation.Nonnull;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -167,30 +169,67 @@ public <R> PartiallyOrderedSet<R> mapAll(@Nonnull final Function<Iterable<? exte
 
     @Nonnull
     public <R> PartiallyOrderedSet<R> mapAll(@Nonnull final Map<T, R> map) {
-        final var mappedElements = Sets.newLinkedHashSet(map.values());
-
-        final var resultDependencyMapBuilder = ImmutableSetMultimap.<R, R>builder();
-        for (final var entry : getTransitiveClosure().entries()) {
-            final var key = entry.getKey();
-            final var value = entry.getValue();
-
-            if (map.containsKey(key) && map.containsKey(value)) {
-                resultDependencyMapBuilder.put(map.get(key), map.get(value));
-            } else {
-                if (!map.containsKey(value)) {
-                    // if key depends on value that does not exist -- do not insert the dependency and also remove key
-                    final var mappedKey = map.get(key);
-                    if (mappedKey != null) {
-                        mappedElements.remove(mappedKey);
+        final var allRemovedBuilder = ImmutableSet.<T>builder();
+        final var leftToRemove = new HashSet<>(Sets.difference(set, map.keySet()));
+
+        final var degreeMap = dependencyMap.entries().stream()
+                .collect(Collectors.groupingBy(
+                        Map.Entry::getKey,
+                        HashMap::new,
+                        Collectors.summingInt(e -> 1)));
+
+        while (!leftToRemove.isEmpty()) {
+            final T toRemove = leftToRemove.iterator().next();
+            for (final var entry: dependencyMap.entries()) {
+                final var key = entry.getKey();
+                if (entry.getValue().equals(toRemove) && degreeMap.containsKey(key)) {
+                    final var updatedDegree = degreeMap.get(key) - 1;
+                    if (updatedDegree == 0) {
+                        leftToRemove.add(key);
+                        degreeMap.remove(key);
+                    } else {
+                        degreeMap.put(key, updatedDegree);
                     }
                 }
             }
+            leftToRemove.remove(toRemove);
+            allRemovedBuilder.add(toRemove);
         }
+        final var allRemovedElements = allRemovedBuilder.build();
+
+        final var mappedElements = set.stream()
+                .filter(value -> !allRemovedElements.contains(value))
+                .map(map::get)
+                .collect(Collectors.toSet());
+        final var resultDependencyMapBuilder = ImmutableSetMultimap.<R, R>builder();
+        dependencyMap.entries().stream()
+                .filter(entry -> !allRemovedElements.contains(entry.getKey()) && !allRemovedElements.contains(entry.getValue()))
+                .map(entry -> Map.entry(map.get(entry.getKey()), map.get(entry.getValue())))
+                .forEach(resultDependencyMapBuilder::put);
 
         // this needs the dependency map to be cleansed (which is done in the constructor)
         return PartiallyOrderedSet.of(mappedElements, resultDependencyMapBuilder.build());
     }
 
+    @Nonnull
+    public <R> PartiallyOrderedSet<R> mapAll(@Nonnull final Multimap<T, R> map) {
+        final var identityMapped = this.filterElements(map::containsKey);
+
+        final var resultSet = identityMapped.getSet().stream()
+                .flatMap(value -> map.get(value).stream())
+                .collect(Collectors.toSet());
+
+        final var resultDependencyMapBuilder = ImmutableSetMultimap.<R, R>builder();
+        for (final var entry: identityMapped.dependencyMap.entries()) {
+            Verify.verify(map.containsKey(entry.getKey()) && map.containsKey(entry.getValue()));
+            map.get(entry.getKey()).forEach(mappedKey -> resultDependencyMapBuilder.putAll(mappedKey, map.get(entry.getValue())));
+            resultSet.addAll(map.get(entry.getKey()));
+        }
+
+        // this needs the dependency map to be cleansed (which is done in the constructor)
+        return PartiallyOrderedSet.of(resultSet, resultDependencyMapBuilder.build());
+    }
+
     /**
      * Method that computes a new partially-ordered set that only retains elements that pass the filtering predicate.
      * The filtering is applied in a way that we do not break dependencies.

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/AggregateIndexExpansionVisitor.java

@@ -53,6 +53,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 
@@ -246,7 +247,7 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
                 throw new RecordCoreException("could not pull grouped value " + groupedValue)
                         .addLogInfo(LogMessageKeys.VALUE, groupedValue);
             }
-            argument = result.get(groupedValue);
+            argument = Iterables.getOnlyElement(result.get(groupedValue));
         } else {
             throw new RecordCoreException("unable to plan group by with non-field value")
                     .addLogInfo(LogMessageKeys.VALUE, groupedValue);
@@ -272,7 +273,7 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
                 throw new RecordCoreException("could not pull grouping value " + groupingValue)
                         .addLogInfo(LogMessageKeys.VALUE, groupingValue);
             }
-            return pulledUpGroupingValuesMap.get(groupingValue);
+            return Iterables.getOnlyElement(pulledUpGroupingValuesMap.get(groupingValue));
         }).collect(ImmutableList.toImmutableList());
 
         final var groupingColsValue = RecordConstructorValue.ofUnnamed(pulledUpGroupingValues);

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/BitmapAggregateIndexExpansionVisitor.java

@@ -41,6 +41,7 @@
 import com.google.common.base.Verify;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 
 import javax.annotation.Nonnull;
@@ -82,7 +83,7 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
                 throw new RecordCoreException("could not pull grouped value " + groupedValue)
                         .addLogInfo(LogMessageKeys.VALUE, groupedValue);
             }
-            argument = result.get(groupedValue);
+            argument = Iterables.getOnlyElement(result.get(groupedValue));
         } else {
             throw new UnsupportedOperationException("unable to plan group by with non-field value " + groupedValue);
         }
@@ -114,7 +115,7 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
                 throw new RecordCoreException("could not pull grouping value " + groupingValue)
                         .addLogInfo(LogMessageKeys.VALUE, groupingValue);
             }
-            return pulledUpGroupingValuesMap.get(groupingValue);
+            return Iterables.getOnlyElement(pulledUpGroupingValuesMap.get(groupingValue));
         }).collect(ImmutableList.toImmutableList());
 
         final var pulledUpGroupingValues = ImmutableList.<Value>builder().addAll(explicitPulledUpGroupingValues).add(implicitGroupingValue).build();

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/KeyExpressionExpansionVisitor.java

@@ -44,6 +44,7 @@
 import com.google.common.base.Verify;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 
 import javax.annotation.Nonnull;
@@ -349,7 +350,7 @@ private static ImmutableList<Placeholder> pullUpPlaceholders(@Nonnull final Grap
                     final var pulledUpValue = pulledUpPlaceholderValuesMap.get(value);
                     final var parameterAlias =
                             Objects.requireNonNull(childExpansionPlaceholderValuesMap.get(value));
-                    return Placeholder.newInstanceWithoutRanges(pulledUpValue, parameterAlias);
+                    return Placeholder.newInstanceWithoutRanges(Iterables.getOnlyElement(pulledUpValue), parameterAlias);
                 })
                 .collect(ImmutableList.toImmutableList());
     }
@@ -372,7 +373,7 @@ private static ImmutableList<Column<? extends Value>> pullUpResultColumns(@Nonnu
                         throw new RecordCoreException("could not pull expansion value " + value)
                                 .addLogInfo(LogMessageKeys.VALUE, value);
                     }
-                    return pulledUpValuesMap.get(value);
+                    return Iterables.getOnlyElement(pulledUpValuesMap.get(value));
                 })
                 .map(Column::unnamedOf)
                 .collect(ImmutableList.toImmutableList());

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/MatchInfo.java

@@ -121,8 +121,8 @@ static ImmutableBiMap<Value, Value> adjustMatchedValueMap(@Nonnull final Correla
                                     candidateLowerExpression.getCorrelatedTo()),
                             candidateAlias);
             final var pulledUpCandidateAggregateValue = candidatePullUpMap.get(candidateAggregateValue);
-            if (pulledUpCandidateAggregateValue != null) {
-                adjustedMatchedAggregateMapBuilder.put(queryAggregateValue, pulledUpCandidateAggregateValue);
+            if (!pulledUpCandidateAggregateValue.isEmpty()) {
+                adjustedMatchedAggregateMapBuilder.put(queryAggregateValue, Iterables.getOnlyElement(pulledUpCandidateAggregateValue));
             }
         }
         return adjustedMatchedAggregateMapBuilder.build();
@@ -506,10 +506,10 @@ public static GroupByMappings pullUpGroupByMappings(@Nonnull final PartialMatch
                                 Sets.difference(queryAggregateValue.getCorrelatedToWithoutChildren(),
                                         queryExpression.getCorrelatedTo()), queryAlias);
                 final var pulledUpQueryAggregateValue = pullUpMap.get(queryAggregateValue);
-                if (pulledUpQueryAggregateValue == null) {
+                if (pulledUpQueryAggregateValue.isEmpty()) {
                     return GroupByMappings.empty();
                 }
-                unmatchedAggregateMapBuilder.put(unmatchedAggregateMapEntry.getKey(), pulledUpQueryAggregateValue);
+                unmatchedAggregateMapBuilder.put(unmatchedAggregateMapEntry.getKey(), Iterables.getOnlyElement(pulledUpQueryAggregateValue));
             }
 
             return GroupByMappings.of(matchedGroupingsMap, matchedAggregatesMap, unmatchedAggregateMapBuilder.build());
@@ -527,8 +527,8 @@ private static ImmutableBiMap<Value, Value> pullUpMatchedValueMap(@Nonnull final
                 final var pullUpMap =
                         queryResultValue.pullUp(ImmutableList.of(queryValue), EvaluationContext.empty(),
                                 AliasMap.emptyMap(), constantAliases, queryAlias);
-                final Value pulledUpQueryValue = pullUpMap.get(queryValue);
-                if (pulledUpQueryValue == null) {
+                final var pulledUpQueryValue = pullUpMap.get(queryValue);
+                if (pulledUpQueryValue.isEmpty()) {
                     continue;
                 }
 
@@ -544,10 +544,10 @@ private static ImmutableBiMap<Value, Value> pullUpMatchedValueMap(@Nonnull final
                                         candidateLowerExpression.getCorrelatedTo()),
                                 candidateAlias);
                 final var pulledUpCandidateAggregateValue = candidatePullUpMap.get(candidateAggregateValue);
-                if (pulledUpCandidateAggregateValue == null) {
+                if (pulledUpCandidateAggregateValue.isEmpty()) {
                     continue;
                 }
-                matchedAggregatesMapBuilder.put(pulledUpQueryValue, pulledUpCandidateAggregateValue);
+                matchedAggregatesMapBuilder.put(Iterables.getOnlyElement(pulledUpQueryValue), Iterables.getOnlyElement(pulledUpCandidateAggregateValue));
             }
             return matchedAggregatesMapBuilder.build();
         }