test fixes

Author: prakharjain09

sql/core/src/main/scala/org/apache/spark/sql/execution/CollapseAggregates.scala

@@ -43,7 +43,7 @@ object CollapseAggregates extends Rule[SparkPlan] {
         if checkIfAggregatesCanBeCollapsed(parent, child) =>
         val completeAggregateExpressions = child.aggregateExpressions.map(_.copy(mode = Complete))
         HashAggregateExec(
-          requiredChildDistributionExpressions = None,
+          requiredChildDistributionExpressions = parent.requiredChildDistributionExpressions,
           groupingExpressions = child.groupingExpressions,
           aggregateExpressions = completeAggregateExpressions,
           aggregateAttributes = completeAggregateExpressions.map(_.resultAttribute),
@@ -55,7 +55,7 @@ object CollapseAggregates extends Rule[SparkPlan] {
         if checkIfAggregatesCanBeCollapsed(parent, child) =>
         val completeAggregateExpressions = child.aggregateExpressions.map(_.copy(mode = Complete))
         SortAggregateExec(
-          requiredChildDistributionExpressions = None,
+          requiredChildDistributionExpressions = parent.requiredChildDistributionExpressions,
           groupingExpressions = child.groupingExpressions,
           aggregateExpressions = completeAggregateExpressions,
           aggregateAttributes = completeAggregateExpressions.map(_.resultAttribute),
@@ -67,7 +67,7 @@ object CollapseAggregates extends Rule[SparkPlan] {
         if checkIfAggregatesCanBeCollapsed(parent, child) =>
         val completeAggregateExpressions = child.aggregateExpressions.map(_.copy(mode = Complete))
         ObjectHashAggregateExec(
-          requiredChildDistributionExpressions = None,
+          requiredChildDistributionExpressions = parent.requiredChildDistributionExpressions,
           groupingExpressions = child.groupingExpressions,
           aggregateExpressions = completeAggregateExpressions,
           aggregateAttributes = completeAggregateExpressions.map(_.resultAttribute),

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDAFSuite.scala

@@ -98,7 +98,8 @@ class HiveUDAFSuite extends QueryTest
 
   test("SPARK-24935: customized Hive UDAF with two aggregation buffers") {
     withTempView("v") {
-      spark.range(100).createTempView("v")
+      // Setting numPartitions > 1 explicitly so that we get two Physical aggregation nodes
+      spark.range(0, 100, 1, 10).createTempView("v")
       val df = sql("SELECT id % 2, mock2(id) FROM v GROUP BY id % 2")
 
       val aggs = collect(df.queryExecution.executedPlan) {