[Enhancement][doris-future] Unify regr_sxx/syy/sxy on AggregateFunctionRegrData

be/src/vec/aggregate_functions/aggregate_function_regr_union.cpp

@@ -60,5 +60,8 @@ void register_aggregate_function_regr_union(AggregateFunctionSimpleFactory& fact
     factory.register_function_both("regr_slope", create_aggregate_function_regr<RegrSlopeFunc>);
     factory.register_function_both("regr_intercept",
                                    create_aggregate_function_regr<RegrInterceptFunc>);
+    factory.register_function_both("regr_sxx", create_aggregate_function_regr<RegrSxxFunc>);
+    factory.register_function_both("regr_syy", create_aggregate_function_regr<RegrSyyFunc>);
+    factory.register_function_both("regr_sxy", create_aggregate_function_regr<RegrSxyFunc>);
 }
 } // namespace doris::vectorized

be/src/vec/aggregate_functions/aggregate_function_regr_union.h

@@ -267,13 +267,48 @@ struct RegrInterceptFunc : AggregateFunctionRegrData<T, true, 2, 2> {
     }
 };
 
+template <PrimitiveType T>
+struct RegrSxxFunc : AggregateFunctionRegrData<T, false, 2, 0> {
+    static constexpr const char* name = "regr_sxx";
+
+    Float64 get_result() const {
+        if (this->n < 1) {
+            return std::numeric_limits<Float64>::quiet_NaN();
+        }
+        return this->sxx();
+    }
+};
+
+template <PrimitiveType T>
+struct RegrSyyFunc : AggregateFunctionRegrData<T, false, 0, 2> {
+    static constexpr const char* name = "regr_syy";
+
+    Float64 get_result() const {
+        if (this->n < 1) {
+            return std::numeric_limits<Float64>::quiet_NaN();
+        }
+        return this->syy();
+    }
+};
+
+template <PrimitiveType T>
+struct RegrSxyFunc : AggregateFunctionRegrData<T, true, 1, 1> {
+    static constexpr const char* name = "regr_sxy";
+
+    Float64 get_result() const {
+        if (this->n < 1) {
+            return std::numeric_limits<Float64>::quiet_NaN();
+        }
+        return this->sxy();
+    }
+};
+
 template <typename RegrFunc, bool y_nullable, bool x_nullable>
 class AggregateFunctionRegrSimple
         : public IAggregateFunctionDataHelper<
                   RegrFunc, AggregateFunctionRegrSimple<RegrFunc, y_nullable, x_nullable>> {
 public:
-    using XInputCol = typename PrimitiveTypeTraits<RegrFunc::Type>::ColumnType;
-    using YInputCol = XInputCol;
+    using InputCol = typename PrimitiveTypeTraits<RegrFunc::Type>::ColumnType;
     using ResultCol = ColumnFloat64;
 
     explicit AggregateFunctionRegrSimple(const DataTypes& argument_types_)
@@ -291,39 +326,20 @@ class AggregateFunctionRegrSimple
 
     void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
              Arena&) const override {
-        bool y_null = false;
-        bool x_null = false;
-        const YInputCol* y_nested_column = nullptr;
-        const XInputCol* x_nested_column = nullptr;
-
+        const auto* y_col = nested_or_null<y_nullable>(columns[0], row_num);
         if constexpr (y_nullable) {
-            const auto& y_column_nullable =
-                    assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*columns[0]);
-            y_null = y_column_nullable.is_null_at(row_num);
-            y_nested_column = assert_cast<const YInputCol*, TypeCheckOnRelease::DISABLE>(
-                    y_column_nullable.get_nested_column_ptr().get());
-        } else {
-            y_nested_column = assert_cast<const YInputCol*, TypeCheckOnRelease::DISABLE>(
-                    (*columns[0]).get_ptr().get());
+            if (y_col == nullptr) {
+                return;
+            }
         }
-
+        const auto* x_col = nested_or_null<x_nullable>(columns[1], row_num);
         if constexpr (x_nullable) {
-            const auto& x_column_nullable =
-                    assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*columns[1]);
-            x_null = x_column_nullable.is_null_at(row_num);
-            x_nested_column = assert_cast<const XInputCol*, TypeCheckOnRelease::DISABLE>(
-                    x_column_nullable.get_nested_column_ptr().get());
-        } else {
-            x_nested_column = assert_cast<const XInputCol*, TypeCheckOnRelease::DISABLE>(
-                    (*columns[1]).get_ptr().get());
-        }
-
-        if (x_null || y_null) {
-            return;
+            if (x_col == nullptr) {
+                return;
+            }
         }
 
-        this->data(place).add(y_nested_column->get_data()[row_num],
-                              x_nested_column->get_data()[row_num]);
+        this->data(place).add(y_col->get_data()[row_num], x_col->get_data()[row_num]);
     }
 
     void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
@@ -355,6 +371,21 @@ class AggregateFunctionRegrSimple
             dst_column.get_data().push_back(result);
         }
     }
+
+private:
+    template <bool Nullable>
+    static ALWAYS_INLINE const InputCol* nested_or_null(const IColumn* col, ssize_t row_num) {
+        if constexpr (Nullable) {
+            const auto& c = assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*col);
+            if (c.is_null_at(row_num)) {
+                return nullptr;
+            }
+            return assert_cast<const InputCol*, TypeCheckOnRelease::DISABLE>(
+                    c.get_nested_column_ptr().get());
+        } else {
+            return assert_cast<const InputCol*, TypeCheckOnRelease::DISABLE>(col->get_ptr().get());
+        }
+    }
 };
 } // namespace doris::vectorized
 

fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java

@@ -77,6 +77,9 @@
 import org.apache.doris.nereids.trees.expressions.functions.agg.QuantileUnion;
 import org.apache.doris.nereids.trees.expressions.functions.agg.RegrIntercept;
 import org.apache.doris.nereids.trees.expressions.functions.agg.RegrSlope;
+import org.apache.doris.nereids.trees.expressions.functions.agg.RegrSxx;
+import org.apache.doris.nereids.trees.expressions.functions.agg.RegrSxy;
+import org.apache.doris.nereids.trees.expressions.functions.agg.RegrSyy;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Retention;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Sem;
 import org.apache.doris.nereids.trees.expressions.functions.agg.SequenceCount;
@@ -173,6 +176,9 @@ private BuiltinAggregateFunctions() {
                 agg(QuantileUnion.class, "quantile_union"),
                 agg(RegrIntercept.class, "regr_intercept"),
                 agg(RegrSlope.class, "regr_slope"),
+                agg(RegrSxx.class, "regr_sxx"),
+                agg(RegrSxy.class, "regr_sxy"),
+                agg(RegrSyy.class, "regr_syy"),
                 agg(Retention.class, "retention"),
                 agg(Sem.class, "sem"),
                 agg(SequenceCount.class, "sequence_count"),

fe/fe-core/src/main/java/org/apache/doris/catalog/FunctionSet.java

@@ -104,6 +104,12 @@ public boolean isNullResultWithOneNullParamFunctions(String funcName) {
 
     public static final String REGR_SLOPE = "regr_slope";
 
+    public static final String REGR_SXX = "regr_sxx";
+
+    public static final String REGR_SXY = "regr_sxy";
+
+    public static final String REGR_SYY = "regr_syy";
+
     public static final String SEQUENCE_COUNT = "sequence_count";
 
     public static final String GROUP_ARRAY_INTERSECT = "group_array_intersect";

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/RegrSxx.java

@@ -0,0 +1,85 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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 org.apache.doris.nereids.trees.expressions.functions.agg;
+
+import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.catalog.Type;
+import org.apache.doris.nereids.exceptions.AnalysisException;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
+import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
+import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression;
+import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
+import org.apache.doris.nereids.types.DataType;
+import org.apache.doris.nereids.types.DoubleType;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+/** regr_sxx agg function. */
+public class RegrSxx extends AggregateFunction
+        implements BinaryExpression, ExplicitlyCastableSignature, AlwaysNullable {
+
+    public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
+            FunctionSignature.ret(DoubleType.INSTANCE).args(DoubleType.INSTANCE, DoubleType.INSTANCE));
+
+    public RegrSxx(Expression arg0, Expression arg1) {
+        this(false, arg0, arg1);
+    }
+
+    public RegrSxx(boolean distinct, Expression arg0, Expression arg1) {
+        super("regr_sxx", distinct, arg0, arg1);
+    }
+
+    public RegrSxx(AggregateFunctionParams functionParams) {
+        super(functionParams);
+    }
+
+    @Override
+    public void checkLegalityBeforeTypeCoercion() {
+        DataType yType = left().getDataType();
+        DataType xType = right().getDataType();
+        if (yType.isOnlyMetricType() || xType.isOnlyMetricType()) {
+            throw new AnalysisException(Type.OnlyMetricTypeErrorMsg);
+        }
+        if (!yType.isNumericType() && !yType.isNullType()) {
+            throw new AnalysisException("regr_sxx requires numeric for first parameter: " + toSql());
+        }
+        if (!xType.isNumericType() && !xType.isNullType()) {
+            throw new AnalysisException("regr_sxx requires numeric for second parameter: " + toSql());
+        }
+    }
+
+    @Override
+    public RegrSxx withDistinctAndChildren(boolean distinct, List<Expression> children) {
+        Preconditions.checkArgument(children.size() == 2);
+        return new RegrSxx(getFunctionParams(distinct, children));
+    }
+
+    @Override
+    public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
+        return visitor.visitRegrSxx(this, context);
+    }
+
+    @Override
+    public List<FunctionSignature> getSignatures() {
+        return SIGNATURES;
+    }
+}

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/RegrSxy.java

@@ -0,0 +1,85 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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 org.apache.doris.nereids.trees.expressions.functions.agg;
+
+import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.catalog.Type;
+import org.apache.doris.nereids.exceptions.AnalysisException;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
+import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
+import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression;
+import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
+import org.apache.doris.nereids.types.DataType;
+import org.apache.doris.nereids.types.DoubleType;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+/** regr_sxy agg function. */
+public class RegrSxy extends AggregateFunction
+        implements BinaryExpression, ExplicitlyCastableSignature, AlwaysNullable {
+
+    public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
+            FunctionSignature.ret(DoubleType.INSTANCE).args(DoubleType.INSTANCE, DoubleType.INSTANCE));
+
+    public RegrSxy(Expression arg0, Expression arg1) {
+        this(false, arg0, arg1);
+    }
+
+    public RegrSxy(boolean distinct, Expression arg0, Expression arg1) {
+        super("regr_sxy", distinct, arg0, arg1);
+    }
+
+    public RegrSxy(AggregateFunctionParams functionParams) {
+        super(functionParams);
+    }
+
+    @Override
+    public void checkLegalityBeforeTypeCoercion() {
+        DataType yType = left().getDataType();
+        DataType xType = right().getDataType();
+        if (yType.isOnlyMetricType() || xType.isOnlyMetricType()) {
+            throw new AnalysisException(Type.OnlyMetricTypeErrorMsg);
+        }
+        if (!yType.isNumericType() && !yType.isNullType()) {
+            throw new AnalysisException("regr_sxy requires numeric for first parameter: " + toSql());
+        }
+        if (!xType.isNumericType() && !xType.isNullType()) {
+            throw new AnalysisException("regr_sxy requires numeric for second parameter: " + toSql());
+        }
+    }
+
+    @Override
+    public RegrSxy withDistinctAndChildren(boolean distinct, List<Expression> children) {
+        Preconditions.checkArgument(children.size() == 2);
+        return new RegrSxy(getFunctionParams(distinct, children));
+    }
+
+    @Override
+    public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
+        return visitor.visitRegrSxy(this, context);
+    }
+
+    @Override
+    public List<FunctionSignature> getSignatures() {
+        return SIGNATURES;
+    }
+}