[presto] Move out M2Y from RegressionState for regr_slope and regr_intercept functions (#25475)

presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java

@@ -68,6 +68,7 @@
 import com.facebook.presto.operator.aggregation.DoubleCovarianceAggregation;
 import com.facebook.presto.operator.aggregation.DoubleHistogramAggregation;
 import com.facebook.presto.operator.aggregation.DoubleRegressionAggregation;
+import com.facebook.presto.operator.aggregation.DoubleRegressionExtendedAggregation;
 import com.facebook.presto.operator.aggregation.DoubleSumAggregation;
 import com.facebook.presto.operator.aggregation.EntropyAggregation;
 import com.facebook.presto.operator.aggregation.GeometricMeanAggregations;
@@ -85,6 +86,7 @@
 import com.facebook.presto.operator.aggregation.RealGeometricMeanAggregations;
 import com.facebook.presto.operator.aggregation.RealHistogramAggregation;
 import com.facebook.presto.operator.aggregation.RealRegressionAggregation;
+import com.facebook.presto.operator.aggregation.RealRegressionExtendedAggregation;
 import com.facebook.presto.operator.aggregation.RealSumAggregation;
 import com.facebook.presto.operator.aggregation.ReduceAggregationFunction;
 import com.facebook.presto.operator.aggregation.SumDataSizeForStats;
@@ -742,7 +744,9 @@ private List<? extends SqlFunction> getBuiltInFunctions(FunctionsConfig function
                 .aggregates(DoubleCovarianceAggregation.class)
                 .aggregates(RealCovarianceAggregation.class)
                 .aggregates(DoubleRegressionAggregation.class)
+                .aggregates(DoubleRegressionExtendedAggregation.class)
                 .aggregates(RealRegressionAggregation.class)
+                .aggregates(RealRegressionExtendedAggregation.class)
                 .aggregates(DoubleCorrelationAggregation.class)
                 .aggregates(RealCorrelationAggregation.class)
                 .aggregates(BitwiseOrAggregation.class)

presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java

@@ -22,6 +22,7 @@
 import com.facebook.presto.operator.aggregation.state.CentralMomentsState;
 import com.facebook.presto.operator.aggregation.state.CorrelationState;
 import com.facebook.presto.operator.aggregation.state.CovarianceState;
+import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState;
 import com.facebook.presto.operator.aggregation.state.RegressionState;
 import com.facebook.presto.operator.aggregation.state.VarianceState;
 import com.facebook.presto.spi.function.AggregationFunctionImplementation;
@@ -145,9 +146,14 @@ public static double getCorrelation(CorrelationState state)
     public static void updateRegressionState(RegressionState state, double x, double y)
     {
         double oldMeanX = state.getMeanX();
-        double oldMeanY = state.getMeanY();
         updateCovarianceState(state, x, y);
         state.setM2X(state.getM2X() + (x - oldMeanX) * (x - state.getMeanX()));
+    }
+
+    public static void updateExtendedRegressionState(ExtendedRegressionState state, double x, double y)
+    {
+        double oldMeanY = state.getMeanY();
+        updateRegressionState(state, x, y);
         state.setM2Y(state.getM2Y() + (y - oldMeanY) * (y - state.getMeanY()));
     }
 
@@ -189,12 +195,12 @@ public static double getRegressionSxy(RegressionState state)
         return state.getC2();
     }
 
-    public static double getRegressionSyy(RegressionState state)
+    public static double getRegressionSyy(ExtendedRegressionState state)
     {
         return state.getM2Y();
     }
 
-    public static double getRegressionR2(RegressionState state)
+    public static double getRegressionR2(ExtendedRegressionState state)
     {
         if (state.getM2X() != 0 && state.getM2Y() == 0) {
             return 1.0;
@@ -311,10 +317,21 @@ public static void mergeRegressionState(RegressionState state, RegressionState o
         long na = state.getCount();
         long nb = otherState.getCount();
         state.setM2X(state.getM2X() + otherState.getM2X() + na * nb * Math.pow(state.getMeanX() - otherState.getMeanX(), 2) / (double) (na + nb));
-        state.setM2Y(state.getM2Y() + otherState.getM2Y() + na * nb * Math.pow(state.getMeanY() - otherState.getMeanY(), 2) / (double) (na + nb));
         updateCovarianceState(state, otherState);
     }
 
+    public static void mergeExtendedRegressionState(ExtendedRegressionState state, ExtendedRegressionState otherState)
+    {
+        if (otherState.getCount() == 0) {
+            return;
+        }
+
+        long na = state.getCount();
+        long nb = otherState.getCount();
+        state.setM2Y(state.getM2Y() + otherState.getM2Y() + na * nb * Math.pow(state.getMeanY() - otherState.getMeanY(), 2) / (double) (na + nb));
+        mergeRegressionState(state, otherState);
+    }
+
     public static String generateAggregationName(String baseName, TypeSignature outputType, List<TypeSignature> inputTypes)
     {
         StringBuilder sb = new StringBuilder();

presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java

@@ -24,15 +24,8 @@
 import com.facebook.presto.spi.function.SqlType;
 
 import static com.facebook.presto.common.type.DoubleType.DOUBLE;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionIntercept;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSlope;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.mergeRegressionState;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.updateRegressionState;
 
@@ -78,100 +71,4 @@ public static void regrIntercept(@AggregationState RegressionState state, BlockB
             out.appendNull();
         }
     }
-
-    @AggregationFunction("regr_sxy")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrSxy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSxy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_sxx")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrSxx(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSxx(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_syy")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrSyy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSyy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_r2")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrR2(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionR2(state);
-        if (Double.isFinite(result)) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_count")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrCount(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionCount(state);
-        if (Double.isFinite(result) && result > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_avgy")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrAvgy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionAvgy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_avgx")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrAvgx(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionAvgx(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
 }

presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java

@@ -0,0 +1,149 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.operator.aggregation;
+
+import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.type.StandardTypes;
+import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState;
+import com.facebook.presto.spi.function.AggregationFunction;
+import com.facebook.presto.spi.function.AggregationState;
+import com.facebook.presto.spi.function.CombineFunction;
+import com.facebook.presto.spi.function.InputFunction;
+import com.facebook.presto.spi.function.OutputFunction;
+import com.facebook.presto.spi.function.SqlType;
+
+import static com.facebook.presto.common.type.DoubleType.DOUBLE;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.mergeExtendedRegressionState;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.updateExtendedRegressionState;
+
+@AggregationFunction
+public class DoubleRegressionExtendedAggregation
+{
+    private DoubleRegressionExtendedAggregation() {}
+
+    @InputFunction
+    public static void input(@AggregationState ExtendedRegressionState state, @SqlType(StandardTypes.DOUBLE) double dependentValue, @SqlType(StandardTypes.DOUBLE) double independentValue)
+    {
+        updateExtendedRegressionState(state, independentValue, dependentValue);
+    }
+
+    @CombineFunction
+    public static void combine(@AggregationState ExtendedRegressionState state, @AggregationState ExtendedRegressionState otherState)
+    {
+        mergeExtendedRegressionState(state, otherState);
+    }
+
+    @AggregationFunction("regr_sxy")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrSxy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_sxx")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrSxx(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxx(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_syy")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrSyy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSyy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_r2")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrR2(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionR2(state);
+        if (Double.isFinite(result)) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_count")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrCount(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionCount(state);
+        if (Double.isFinite(result) && result > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgy")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrAvgy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgx")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrAvgx(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgx(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+}