Fix MAP_UNION aggregate function to support ORDERY BY clause

prestodb · Apr 23, 2021 · 2cfbf87 · 2cfbf87
1 parent 3c8c4ec
commit 2cfbf87
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/presto-main/src/main/java/com/facebook/presto/operator/aggregation/MapUnionAggregation.java b/presto-main/src/main/java/com/facebook/presto/operator/aggregation/MapUnionAggregation.java
@@ -93,7 +93,7 @@ private static InternalAggregationFunction generateAggregation(Type keyType, Typ
                 outputType);
 
         GenericAccumulatorFactoryBinder factory = AccumulatorCompiler.generateAccumulatorFactoryBinder(metadata, classLoader);
-        return new InternalAggregationFunction(NAME, inputTypes, ImmutableList.of(intermediateType), outputType, true, false, factory);
+        return new InternalAggregationFunction(NAME, inputTypes, ImmutableList.of(intermediateType), outputType, true, true, factory);
     }
 
     private static List<ParameterMetadata> createInputParameterMetadata(Type inputType)

diff --git a/presto-main/src/test/java/com/facebook/presto/sql/query/TestOrderedAggregation.java b/presto-main/src/test/java/com/facebook/presto/sql/query/TestOrderedAggregation.java
@@ -122,6 +122,22 @@ public void testAggregationWithOrderBy()
         assertions.assertQuery(
                 "SELECT multimap_agg(x, y ORDER BY z) FROM (VALUES (1, 2, 2), (1, 5, 5), (2, 1, 5), (3, 4, 4), (2, 5, 1), (1, 1, 1)) t(x, y, z)",
                 "VALUES map_from_entries(ARRAY[row(1, ARRAY[1, 2, 5]), row(2, ARRAY[5, 1]), row(3, ARRAY[4])])");
+
+        assertions.assertQuery(
+                "SELECT map_union(x ORDER BY y) FROM (VALUES (map(array[1], array['b']), 1), (map(array[1], array['a']), 10)) t(x,y)",
+                "VALUES map_from_entries(ARRAY[(1,'b')])");
+
+        assertions.assertQuery(
+                "SELECT map_union(x ORDER BY y desc) FROM (VALUES (map(array[1], array['b']), 1), (map(array[1], array['a']), 10)) t(x,y)",
+                "VALUES map_from_entries(ARRAY[(1,'a')])");
+
+        assertions.assertQuery(
+                "SELECT map_union(x ORDER BY z) FROM (VALUES (map(array[1], array['c']), 1, 2), (map(array[1], array['b']), 1, 3), (map(array[1], array['a']), 10, 4)) t(x,y,z) GROUP BY y",
+                "VALUES map_from_entries(ARRAY[(1,'c')]), map_from_entries(ARRAY[(1,'a')])");
+
+        assertions.assertQuery(
+                "SELECT map_union(x ORDER BY z desc) FROM (VALUES (map(array[1], array['c']), 1, 2), (map(array[1], array['b']), 1, 3), (map(array[1], array['d']), 1, 4), (map(array[1], array['a']), 10, 4)) t(x,y,z) GROUP BY y",
+                "VALUES map_from_entries(ARRAY[(1,'d')]), map_from_entries(ARRAY[(1,'a')])");
     }
 
     @Test