Fix bug in recursive assignment resolver

The assignments map was not always being updated when a variable was resolved
prestodb · Oct 31, 2023 · d2df443 · d2df443
1 parent 7e89851
commit d2df443
Showing 1 changed file with 5 additions and 10 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ReorderJoins.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ReorderJoins.java
@@ -713,9 +713,7 @@ private static class JoinNodeFlattener
             private Assignments resolveAssignments(Map<VariableReferenceExpression, RowExpression> assignments, Set<VariableReferenceExpression> availableVariables)
             {
                 HashSet<VariableReferenceExpression> resolvedVariables = new HashSet<>();
-                for (VariableReferenceExpression variable : assignments.keySet()) {
-                    resolveVariable(variable, resolvedVariables, assignments, availableVariables);
-                }
+                ImmutableList.copyOf(assignments.keySet()).forEach(variable -> resolveVariable(variable, resolvedVariables, assignments, availableVariables));
 
                 return Assignments.builder().putAll(assignments).build();
             }
@@ -724,17 +722,14 @@ private void resolveVariable(VariableReferenceExpression variable, HashSet<Varia
                     RowExpression> assignments, Set<VariableReferenceExpression> availableVariables)
             {
                 RowExpression expression = assignments.get(variable);
-
                 Sets.SetView<VariableReferenceExpression> variablesToResolve = Sets.difference(Sets.difference(extractUnique(expression), availableVariables), resolvedVariables);
-                if (variablesToResolve.isEmpty()) {
-                    resolvedVariables.add(variable);
-                    return;
-                }
 
+                // Recursively resolve any unresolved variables
                 variablesToResolve.forEach(variableToResolve -> resolveVariable(variableToResolve, resolvedVariables, assignments, availableVariables));
-                // Modify the assignments to replace the variables with the resolved expressions
+
+                // Modify the assignment for the variable : Replace it with the now resolved constituent variables
                 assignments.put(variable, replaceExpression(expression, assignments));
-                // Mark the variable as resolved
+                // Mark this variable as resolved
                 resolvedVariables.add(variable);
             }