From 981ead9032098ede412182a3d44b759a71c2ab3b Mon Sep 17 00:00:00 2001 From: mch_ucchi <41606806+sohardforaname@users.noreply.github.com> Date: Thu, 6 Apr 2023 16:57:04 +0800 Subject: [PATCH] [feature](Nereids) support binary arithmetic function (#18213) support binary arithmetic functions like: add(op1, op2) -> op1 + op2 subtract(op1, op2) -> op1 - op2 multiply(op1, op2) -> op1 * op2 divide(op1, op2) -> op1 / op2 mod(op1, op2) -> op1 % op2 --- .../analysis/ArithmeticFunctionBinder.java | 63 +++++++++++++++++++ .../rules/analysis/FunctionBinder.java | 8 +++ .../sql_functions/test_arith_functions.groovy | 38 +++++++++++ 3 files changed, 109 insertions(+) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ArithmeticFunctionBinder.java create mode 100644 regression-test/suites/nereids_p0/sql_functions/test_arith_functions.groovy diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ArithmeticFunctionBinder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ArithmeticFunctionBinder.java new file mode 100644 index 0000000000..0a1d2a788d --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ArithmeticFunctionBinder.java @@ -0,0 +1,63 @@ +// 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.rules.analysis; + +import org.apache.doris.nereids.trees.expressions.Add; +import org.apache.doris.nereids.trees.expressions.BitAnd; +import org.apache.doris.nereids.trees.expressions.BitNot; +import org.apache.doris.nereids.trees.expressions.BitOr; +import org.apache.doris.nereids.trees.expressions.BitXor; +import org.apache.doris.nereids.trees.expressions.Divide; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.IntegralDivide; +import org.apache.doris.nereids.trees.expressions.Mod; +import org.apache.doris.nereids.trees.expressions.Multiply; +import org.apache.doris.nereids.trees.expressions.Subtract; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; + +import com.google.common.collect.ImmutableMap; + +import java.util.List; +import java.util.Map; + +/** + * bind arithmetic function + */ +public class ArithmeticFunctionBinder { + private static final NullLiteral DUMMY_EXPRESSION = new NullLiteral(); + private static final Map FUNCTION_TO_EXPRESSION = ImmutableMap.builder() + .put("add", new Add(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("subtract", new Subtract(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("multiply", new Multiply(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("divide", new Divide(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("int_divide", new IntegralDivide(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("mod", new Mod(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("bitand", new BitAnd(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("bitor", new BitOr(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("bitxor", new BitXor(DUMMY_EXPRESSION, DUMMY_EXPRESSION)) + .put("bitnot", new BitNot(DUMMY_EXPRESSION)) + .build(); + + public boolean isBinaryArithmetic(String functionName) { + return FUNCTION_TO_EXPRESSION.containsKey(functionName.toLowerCase()); + } + + public Expression bindBinaryArithmetic(String functionName, List children) { + return FUNCTION_TO_EXPRESSION.get(functionName.toLowerCase()).withChildren(children); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/FunctionBinder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/FunctionBinder.java index 031d8c74d7..6e8a98d244 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/FunctionBinder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/FunctionBinder.java @@ -96,6 +96,14 @@ public class FunctionBinder extends DefaultExpressionRewriter { .build() : (List) unboundFunction.getArguments(); + // we will change arithmetic function like add(), subtract(), bitnot() to the corresponding objects rather than + // BoundFunction. + ArithmeticFunctionBinder functionBinder = new ArithmeticFunctionBinder(); + if (functionBinder.isBinaryArithmetic(unboundFunction.getName())) { + return functionBinder.bindBinaryArithmetic(unboundFunction.getName(), unboundFunction.children()) + .accept(this, context); + } + FunctionBuilder builder = functionRegistry.findFunctionBuilder(functionName, arguments); BoundFunction boundFunction = builder.build(functionName, arguments); return TypeCoercionUtils.processBoundFunction(boundFunction); diff --git a/regression-test/suites/nereids_p0/sql_functions/test_arith_functions.groovy b/regression-test/suites/nereids_p0/sql_functions/test_arith_functions.groovy new file mode 100644 index 0000000000..3111a740f9 --- /dev/null +++ b/regression-test/suites/nereids_p0/sql_functions/test_arith_functions.groovy @@ -0,0 +1,38 @@ +// 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. + +suite("test_arith_functions") { + sql "SET enable_nereids_planner=true" + sql "SET enable_fallback_to_original_planner=false" + sql "use test_query_db" + + test { + sql 'select add(1, 1), subtract(1, 1), multiply(2, 2), divide(3.0, 2.0), mod(3.0, 1.3)' + result([[2, 0, 4, 1.5, 0.400000000]]) + } + test { + sql 'select int_divide(1, 1), bitand(1, 1), bitor(2, 2), bitxor(3.0, 2.0), bitnot(3.0)' + result([[1L, 1, 2, 1L, -4L]]) + } + test { + sql 'select add(k1, k2) + subtract(k2, k3) + multiply(k3, k4), cast(divide(k4, k3) + mod(k4, k3) as bigint) from test order by k1 limit 1' + result([[11022916880, 11902L]]) + } +// test { +// sql 'select int_divide(k1, k2), bitand(k2, k3), bitor(k3, k4), bitxor(k4, k3), bitnot(k4) from test order by k1' +// } +}