Add more Linear Regression functions to Presto

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

@@ -145,8 +145,10 @@ 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()));
+        state.setM2Y(state.getM2Y() + (y - oldMeanY) * (y - state.getMeanY()));
     }
 
     public static double getRegressionSlope(RegressionState state)
@@ -167,6 +169,41 @@ public static double getRegressionIntercept(RegressionState state)
         return meanY - slope * meanX;
     }
 
+    public static double getRegressionAvgy(RegressionState state)
+    {
+        return state.getMeanY();
+    }
+
+    public static double getRegressionAvgx(RegressionState state)
+    {
+        return state.getMeanX();
+    }
+
+    public static double getRegressionSxx(RegressionState state)
+    {
+        return state.getM2X();
+    }
+
+    public static double getRegressionSxy(RegressionState state)
+    {
+        return state.getC2();
+    }
+
+    public static double getRegressionSyy(RegressionState state)
+    {
+        return state.getM2Y();
+    }
+
+    public static double getRegressionR2(RegressionState state)
+    {
+        return Math.pow(state.getC2(), 2) / (state.getM2X() * state.getM2Y());
+    }
+
+    public static double getRegressionCount(RegressionState state)
+    {
+        return state.getCount();
+    }
+
     public static void mergeVarianceState(VarianceState state, VarianceState otherState)
     {
         long count = otherState.getCount();
@@ -265,6 +302,7 @@ 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);
     }
 

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

@@ -24,8 +24,15 @@
 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;
 
@@ -71,4 +78,95 @@ 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);
+        if (Double.isFinite(result)) {
+            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);
+        if (Double.isFinite(result)) {
+            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);
+        if (Double.isFinite(result)) {
+            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)) {
+            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);
+        if (Double.isFinite(result)) {
+            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);
+        if (Double.isFinite(result)) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
 }

presto-main/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java

@@ -24,8 +24,15 @@
 import com.facebook.presto.spi.function.SqlType;
 
 import static com.facebook.presto.common.type.RealType.REAL;
+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 java.lang.Float.floatToRawIntBits;
 import static java.lang.Float.intBitsToFloat;
 
@@ -71,4 +78,95 @@ public static void regrIntercept(@AggregationState RegressionState state, BlockB
             out.appendNull();
         }
     }
+
+    @AggregationFunction("regr_sxy")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrSxy(@AggregationState RegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxy(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_sxx")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrSxx(@AggregationState RegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxx(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_syy")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrSyy(@AggregationState RegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSyy(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_r2")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrR2(@AggregationState RegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionR2(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_count")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrCount(@AggregationState RegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionCount(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgy")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrAvgy(@AggregationState RegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgy(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgx")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrAvgx(@AggregationState RegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgx(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
 }

presto-main/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java

@@ -19,4 +19,8 @@ public interface RegressionState
     double getM2X();
 
     void setM2X(double value);
+
+    double getM2Y();
+
+    void setM2Y(double value);
 }

presto-main/src/test/java/com/facebook/presto/operator/aggregation/TestDoubleRegrAvgxAggregation.java

@@ -0,0 +1,79 @@
+/*
+ * 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.Block;
+import com.facebook.presto.common.type.StandardTypes;
+import com.google.common.collect.ImmutableList;
+import org.testng.annotations.Test;
+
+import java.util.List;
+
+import static com.facebook.presto.block.BlockAssertions.createDoubleSequenceBlock;
+import static com.facebook.presto.block.BlockAssertions.createDoublesBlock;
+import static com.google.common.base.Preconditions.checkArgument;
+
+public class TestDoubleRegrAvgxAggregation
+        extends AbstractTestAggregationFunction
+{
+    @Override
+    public Block[] getSequenceBlocks(int start, int length)
+    {
+        return new Block[] {createDoubleSequenceBlock(start, start + length), createDoubleSequenceBlock(start + 2, start + 2 + length)};
+    }
+
+    @Override
+    protected String getFunctionName()
+    {
+        return "regr_avgx";
+    }
+
+    @Override
+    protected List<String> getFunctionParameterTypes()
+    {
+        return ImmutableList.of(StandardTypes.DOUBLE, StandardTypes.DOUBLE);
+    }
+
+    @Override
+    public Object getExpectedValue(int start, int length)
+    {
+        if (length == 0) {
+            return 0.0;
+        }
+
+        double expected = 0.0;
+        for (int i = start; i < start + length; i++) {
+            expected += (i + 2);
+        }
+        return expected / length;
+    }
+
+    @Test
+    public void testNonTrivialResult()
+    {
+        testNonTrivialAggregation(new Double[] {1.0, 2.0, 3.0, 4.0, 5.0}, new Double[] {1.0, 4.0, 9.0, 16.0, 25.0});
+        testNonTrivialAggregation(new Double[] {1.0, 4.0, 9.0, 16.0, 25.0}, new Double[] {1.0, 2.0, 3.0, 4.0, 5.0});
+    }
+
+    private void testNonTrivialAggregation(Double[] y, Double[] x)
+    {
+        double expected = 0.0;
+        for (int i = 0; i < x.length; i++) {
+            expected += x[i];
+        }
+        expected = expected / x.length;
+        checkArgument(Double.isFinite(expected) && expected != 0., "Expected result is trivial");
+        testAggregation(expected, createDoublesBlock(y), createDoublesBlock(x));
+    }
+}