[enhancement](function) floor/ceil/round/round_bankers can use column as scale argument (#34391)

This commit is contained in:
Chester
2024-05-06 22:17:31 +08:00
committed by yiguolei
parent c22f42121b
commit f7900b53ce
13 changed files with 1522 additions and 796 deletions

View File

@ -115,30 +115,8 @@ public class FunctionCallExpr extends Expr {
return returnType;
}
};
java.util.function.BiFunction<ArrayList<Expr>, Type, Type> roundRule = (children, returnType) -> {
Preconditions.checkArgument(children != null && children.size() > 0);
if (children.size() == 1 && children.get(0).getType().isDecimalV3()) {
return ScalarType.createDecimalV3Type(children.get(0).getType().getPrecision(), 0);
} else if (children.size() == 2) {
Preconditions.checkArgument(children.get(1) instanceof IntLiteral
|| (children.get(1) instanceof CastExpr
&& children.get(1).getChild(0) instanceof IntLiteral),
"2nd argument of function round/floor/ceil must be literal");
if (children.get(1) instanceof CastExpr && children.get(1).getChild(0) instanceof IntLiteral) {
children.get(1).getChild(0).setType(children.get(1).getType());
children.set(1, children.get(1).getChild(0));
} else {
children.get(1).setType(Type.INT);
}
int scaleArg = (int) (((IntLiteral) children.get(1)).getValue());
return ScalarType.createDecimalV3Type(children.get(0).getType().getPrecision(),
Math.min(Math.max(scaleArg, 0), ((ScalarType) children.get(0).getType()).decimalScale()));
} else {
return returnType;
}
};
java.util.function.BiFunction<ArrayList<Expr>, Type, Type> truncateRule = (children, returnType) -> {
java.util.function.BiFunction<ArrayList<Expr>, Type, Type> roundRule = (children, returnType) -> {
Preconditions.checkArgument(children != null && children.size() > 0);
if (children.size() == 1 && children.get(0).getType().isDecimalV3()) {
return ScalarType.createDecimalV3Type(children.get(0).getType().getPrecision(), 0);
@ -268,7 +246,7 @@ public class FunctionCallExpr extends Expr {
PRECISION_INFER_RULE.put("dround", roundRule);
PRECISION_INFER_RULE.put("dceil", roundRule);
PRECISION_INFER_RULE.put("dfloor", roundRule);
PRECISION_INFER_RULE.put("truncate", truncateRule);
PRECISION_INFER_RULE.put("truncate", roundRule);
}
public static final ImmutableSet<String> TIME_FUNCTIONS_WITH_PRECISION = new ImmutableSortedSet.Builder(

View File

@ -20,13 +20,10 @@ package org.apache.doris.nereids.trees.expressions.functions;
import org.apache.doris.catalog.FunctionSignature;
import org.apache.doris.nereids.trees.expressions.Cast;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.scalar.Truncate;
import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral;
import org.apache.doris.nereids.types.DecimalV3Type;
import org.apache.doris.nereids.types.coercion.Int32OrLessType;
import com.google.common.base.Preconditions;
/** ComputePrecisionForRound */
public interface ComputePrecisionForRound extends ComputePrecision {
@Override
@ -40,34 +37,19 @@ public interface ComputePrecisionForRound extends ComputePrecision {
Expression floatLength = getArgument(1);
int scale;
if (this instanceof Truncate) {
if (floatLength.isLiteral() || (
floatLength instanceof Cast && floatLength.child(0).isLiteral()
&& floatLength.child(0).getDataType() instanceof Int32OrLessType)) {
// Scale argument is a literal or cast from other literal
if (floatLength instanceof Cast) {
scale = ((IntegerLikeLiteral) floatLength.child(0)).getIntValue();
} else {
scale = ((IntegerLikeLiteral) floatLength).getIntValue();
}
scale = Math.min(Math.max(scale, 0), decimalV3Type.getScale());
} else {
// Truncate could use Column as its scale argument.
// Result scale will always same with input Decimal in this situation.
scale = decimalV3Type.getScale();
}
} else {
Preconditions.checkArgument(floatLength.getDataType() instanceof Int32OrLessType
&& (floatLength.isLiteral() || (
floatLength instanceof Cast && floatLength.child(0).isLiteral()
&& floatLength.child(0).getDataType() instanceof Int32OrLessType)),
"2nd argument of function round/floor/ceil must be literal");
if (floatLength.isLiteral() || (floatLength instanceof Cast && floatLength.child(0).isLiteral()
&& floatLength.child(0).getDataType() instanceof Int32OrLessType)) {
// Scale argument is a literal or cast from other literal
if (floatLength instanceof Cast) {
scale = ((IntegerLikeLiteral) floatLength.child(0)).getIntValue();
} else {
scale = ((IntegerLikeLiteral) floatLength).getIntValue();
}
scale = Math.min(Math.max(scale, 0), decimalV3Type.getScale());
} else {
// Func could use Column as its scale argument.
// Result scale will always same with input Decimal in this situation.
scale = decimalV3Type.getScale();
}
return signature.withArgumentType(0, decimalV3Type)