[GLUTEN-13340][VL] Fix regr_r2's intermediate field order

backends-velox/src/main/scala/org/apache/gluten/utils/VeloxIntermediateData.scala

@@ -36,9 +36,14 @@ object VeloxIntermediateData {
   // Accumulator is mapped to both xAvg and avg in Spark's RegrSlope aggBufferAttributes. In this
   // scenario, when passing the output of Spark's partial aggregation to Velox, we only need to
   // take one of them.
-  // Corr, RegrR2
+  // Corr
   private val veloxCorrIntermediateDataOrder: Seq[Seq[String]] =
     Seq("ck", "n", "xMk", "yMk", "xAvg", "yAvg").map(Seq(_))
+  // RegrR2. Velox's accumulator is (c2, count, m2X, meanX, meanY, m2Y) with X = the second
+  // argument, while Spark's RegrR2(y, x) extends PearsonCorrelation(y, x), so Spark's x* buffer
+  // fields describe y and its y* fields describe x.
+  private val veloxRegrR2IntermediateDataOrder: Seq[Seq[String]] =
+    Seq("ck", "n", "yMk", "yAvg", "xAvg", "xMk").map(Seq(_))
   // CovPopulation, CovSample
   private val veloxCovarIntermediateDataOrder: Seq[Seq[String]] =
     Seq("ck", "n", "xAvg", "yAvg").map(Seq(_))
@@ -90,6 +95,9 @@ object VeloxIntermediateData {
    */
   def veloxIntermediateDataOrder(aggFunc: AggregateFunction): Seq[Seq[String]] = {
     aggFunc match {
+      // RegrR2 extends PearsonCorrelation, it must be placed before PearsonCorrelation.
+      case _: RegrR2 =>
+        veloxRegrR2IntermediateDataOrder
       case _: PearsonCorrelation =>
         veloxCorrIntermediateDataOrder
       case _: CovPopulation | _: CovSample =>

backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala

@@ -20,7 +20,8 @@ import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
 import org.apache.gluten.extension.columnar.validator.FallbackInjects
 
 import org.apache.spark.SparkConf
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge}
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge, RegrR2}
 import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
 import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
 import org.apache.spark.sql.internal.SQLConf
@@ -418,6 +419,69 @@ abstract class VeloxAggregateFunctionsSuite extends VeloxWholeStageTransformerSu
     }
   }
 
+  private def withRegrR2Table(f: => Unit): Unit = {
+    // k = 3 has a constant x and k = 4 a constant y.
+    withTable("regr_r2_src") {
+      Seq(
+        (1, 1.0, 2.0, 1),
+        (1, 2.0, 4.5, 2),
+        (1, 3.5, 3.5, 1),
+        (1, 4.0, 8.0, 3),
+        (2, 1.5, 1.0, 2),
+        (2, 2.5, 3.0, 2),
+        (2, 3.0, 6.5, 4),
+        (2, 5.0, 5.0, 1),
+        (3, 2.0, 1.0, 1),
+        (3, 2.0, 3.0, 2),
+        (3, 2.0, 7.0, 3),
+        (4, 1.0, 5.0, 1),
+        (4, 3.0, 5.0, 2),
+        (4, 6.0, 5.0, 3)
+      ).toDF("k", "x", "y", "z").write.format("parquet").saveAsTable("regr_r2_src")
+      f
+    }
+  }
+
+  private def checkRegrR2RunsInBothEngines(df: DataFrame): Unit = {
+    val regrR2Aggs = getExecutedPlan(df).collect {
+      case agg: BaseAggregateExec
+          if agg.aggregateExpressions.exists(_.aggregateFunction.isInstanceOf[RegrR2]) =>
+        agg
+    }
+    assert(
+      regrR2Aggs.exists(_.isInstanceOf[HashAggregateExecTransformer]),
+      "Expected a Velox aggregation that computes regr_r2")
+    assert(
+      regrR2Aggs.exists(!_.isInstanceOf[HashAggregateExecTransformer]),
+      "Expected a vanilla Spark aggregation that computes regr_r2")
+  }
+
+  test("regr_r2 with a Velox partial and a vanilla Spark final aggregation") {
+    withRegrR2Table {
+      // Spark plans a single DISTINCT as four aggregations that all carry regr_r2.
+      // histogram_numeric is not offloaded, so the last two run in vanilla Spark and merge the
+      // buffer Velox produced.
+      runQueryAndCompare(
+        "select k, regr_r2(y, x), histogram_numeric(distinct z, 3) from regr_r2_src group by k") {
+        checkRegrR2RunsInBothEngines
+      }
+    }
+  }
+
+  test("regr_r2 with a vanilla Spark partial and a Velox final aggregation") {
+    withRegrR2Table {
+      // For a constant x or y, Velox's final regr_r2 result matches Spark only after SPARK-55969,
+      // so only k = 1 and k = 2 are compared here.
+      FallbackInjects.fallbackOn {
+        case agg: BaseAggregateExec => agg.aggregateExpressions.exists(_.mode == Partial)
+      } {
+        runQueryAndCompare("select k, regr_r2(y, x) from regr_r2_src where k <= 2 group by k") {
+          checkRegrR2RunsInBothEngines
+        }
+      }
+    }
+  }
+
   test("regr_slope") {
     runQueryAndCompare("""
                          |select regr_slope(l_partkey, l_suppkey) from lineitem;