Fix failure when validating pre-sorted input properties

Continuation of d5e3a9a to handle the case where multiple constant columns appear in a different order.
trinodb · Oct 30, 2021 · 6507f83 · 6507f83
1 parent e537a09
commit 6507f83
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 24 deletions.
diff --git a/...trino-main/src/main/java/io/trino/sql/planner/sanity/ValidateLimitWithPresortedInput.java b/...trino-main/src/main/java/io/trino/sql/planner/sanity/ValidateLimitWithPresortedInput.java
@@ -31,6 +31,9 @@
 import io.trino.sql.planner.plan.PlanVisitor;
 import io.trino.sql.planner.sanity.PlanSanityChecker.Checker;
 
+import java.util.HashSet;
+import java.util.Set;
+
 import static com.google.common.collect.Iterators.peekingIterator;
 import static io.trino.sql.planner.optimizations.LocalProperties.normalizeAndPrune;
 import static io.trino.sql.planner.optimizations.StreamPropertyDerivations.derivePropertiesRecursively;
@@ -95,33 +98,25 @@ public Void visitLimit(LimitNode node, Void context)
 
             StreamProperties properties = derivePropertiesRecursively(node.getSource(), metadata, typeOperators, session, types, typeAnalyzer);
 
-            // We do not use LocalProperties#match to fully verify sorting properties
-            // as OrderingScheme of input is not tracked by LimitNode
-            PeekingIterator<Symbol> expected = peekingIterator(node.getPreSortedInputs().iterator());
-
-            for (LocalProperty<Symbol> property : normalizeAndPrune(properties.getLocalProperties())) {
-                if (!expected.hasNext()) {
-                    // all properties satisfied
-                    break;
-                }
-
-                Symbol column = expected.peek();
-                if (property instanceof SortingProperty && ((SortingProperty<Symbol>) property).getColumn().equals(column) ||
-                        property instanceof ConstantProperty && ((ConstantProperty<Symbol>) property).getColumn().equals(column)) {
-                    expected.next();
-                    continue;
+            PeekingIterator<LocalProperty<Symbol>> actuals = peekingIterator(normalizeAndPrune(properties.getLocalProperties()).iterator());
+
+            Set<Symbol> satisfied = new HashSet<>();
+            for (Symbol expected : node.getPreSortedInputs()) {
+                while (actuals.hasNext()) {
+                    LocalProperty<Symbol> actual = actuals.peek();
+                    if (actual instanceof ConstantProperty ||
+                            (actual instanceof SortingProperty && ((SortingProperty<Symbol>) actual).getColumn().equals(expected))) {
+                        satisfied.addAll(actual.getColumns());
+                        actuals.next();
+                    }
+                    else {
+                        break;
+                    }
                 }
 
-                if (property instanceof ConstantProperty) {
-                    // unrelated constant columns don't matter for the purpose of satisfying sorting requirements
-                    continue;
+                if (!satisfied.contains(expected)) {
+                    throw new VerifyException(format("Expected Limit input to be sorted by: %s, but was %s", node.getPreSortedInputs(), properties.getLocalProperties()));
                 }
-
-                break;
-            }
-
-            if (expected.hasNext()) {
-                throw new VerifyException(format("Expected Limit input to be sorted by: %s, but was %s", node.getPreSortedInputs(), properties.getLocalProperties()));
             }
 
             return null;

diff --git a/testing/trino-tests/src/test/java/io/trino/tests/AbstractTestEngineOnlyQueries.java b/testing/trino-tests/src/test/java/io/trino/tests/AbstractTestEngineOnlyQueries.java
@@ -6171,6 +6171,14 @@ public void testPartialLimitWithPresortedConstantInputs()
                 "ORDER BY orderkey " +
                 "LIMIT 1"))
                 .matches("VALUES (BIGINT '1', BIGINT '370')");
+
+        assertThat(query("SELECT " +
+                "         'name' as name, " +
+                "         'age' as age " +
+                "         FROM customer " +
+                "         ORDER BY age, name " +
+                "         LIMIT 1"))
+                .matches("VALUES ('name', 'age')");
     }
 
     private static ZonedDateTime zonedDateTime(String value)