This is an automated email from the ASF dual-hosted git repository.
kxiao pushed a commit to branch branch-2.0
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-2.0 by this push:
new e9d944c5f32 [branch-2.0](fix) Fix be crash about mark_first/last_n
#35884 (#36357)
e9d944c5f32 is described below
commit e9d944c5f327424721c07d694387cc864dd2599d
Author: zclllyybb <zhaochan...@selectdb.com>
AuthorDate: Wed Jun 19 12:17:51 2024 +0800
[branch-2.0](fix) Fix be crash about mark_first/last_n #35884 (#36357)
fix crash when mark_first/last_n gets negative 2nd arg input.
and make mark_first/last_n support literal only for 2nd arg
it's a part of https://github.com/apache/doris/pull/35884
---
be/src/vec/functions/function_string.h | 20 ++++++++++----------
.../expressions/functions/scalar/MaskFirstN.java | 8 ++++++++
.../expressions/functions/scalar/MaskLastN.java | 8 ++++++++
.../suites/correctness_p0/test_mask_function.groovy | 18 ++++++++++++++++++
4 files changed, 44 insertions(+), 10 deletions(-)
diff --git a/be/src/vec/functions/function_string.h
b/be/src/vec/functions/function_string.h
index 537b5748f3c..6b7bfca9808 100644
--- a/be/src/vec/functions/function_string.h
+++ b/be/src/vec/functions/function_string.h
@@ -506,28 +506,28 @@ public:
Status execute_impl(FunctionContext* context, Block& block, const
ColumnNumbers& arguments,
size_t result, size_t input_rows_count) override {
- DCHECK_GE(arguments.size(), 1);
- DCHECK_LE(arguments.size(), 2);
-
- int n = -1;
+ int n = -1; // means unassigned
auto res = ColumnString::create();
auto col =
block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
const ColumnString& source_column = assert_cast<const
ColumnString&>(*col);
if (arguments.size() == 2) {
- auto& col = *block.get_by_position(arguments[1]).column;
+ const auto& col = *block.get_by_position(arguments[1]).column;
+ // the 2nd arg is const. checked in fe.
+ if (col.get_int(0) < 0) [[unlikely]] {
+ return Status::InvalidArgument(
+ "function {} only accept non-negative input for 2nd
argument but got {}",
+ name, col.get_int(0));
+ }
n = col.get_int(0);
- } else if (arguments.size() > 2) {
- return Status::InvalidArgument(
- fmt::format("too many arguments for function {}",
get_name()));
}
- if (n == -1) {
+ if (n == -1) { // no 2nd arg, just mask all
FunctionMask::vector_mask(source_column, *res,
FunctionMask::DEFAULT_UPPER_MASK,
FunctionMask::DEFAULT_LOWER_MASK,
FunctionMask::DEFAULT_NUMBER_MASK);
- } else if (n >= 0) {
+ } else { // n >= 0
vector(source_column, n, *res);
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskFirstN.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskFirstN.java
index 81a968067c2..33e19d468e8 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskFirstN.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskFirstN.java
@@ -18,6 +18,7 @@
package org.apache.doris.nereids.trees.expressions.functions.scalar;
import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.expressions.Expression;
import
org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable;
@@ -65,6 +66,13 @@ public class MaskFirstN extends ScalarFunction implements
ExplicitlyCastableSign
return new MaskFirstN(children.get(0), children.get(1));
}
+ @Override
+ public void checkLegalityAfterRewrite() {
+ if (arity() == 2 && !child(1).isLiteral()) {
+ throw new AnalysisException("mask_first_n must accept literal for
2nd argument");
+ }
+ }
+
@Override
public List<FunctionSignature> getSignatures() {
return SIGNATURES;
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskLastN.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskLastN.java
index cb8246f04ab..eafb85ee89b 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskLastN.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MaskLastN.java
@@ -18,6 +18,7 @@
package org.apache.doris.nereids.trees.expressions.functions.scalar;
import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.expressions.Expression;
import
org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable;
@@ -65,6 +66,13 @@ public class MaskLastN extends ScalarFunction implements
ExplicitlyCastableSigna
return new MaskLastN(children.get(0), children.get(1));
}
+ @Override
+ public void checkLegalityAfterRewrite() {
+ if (arity() == 2 && !child(1).isLiteral()) {
+ throw new AnalysisException("mask_last_n must accept literal for
2nd argument");
+ }
+ }
+
@Override
public List<FunctionSignature> getSignatures() {
return SIGNATURES;
diff --git a/regression-test/suites/correctness_p0/test_mask_function.groovy
b/regression-test/suites/correctness_p0/test_mask_function.groovy
index b242e72eccc..e487e51897c 100644
--- a/regression-test/suites/correctness_p0/test_mask_function.groovy
+++ b/regression-test/suites/correctness_p0/test_mask_function.groovy
@@ -16,6 +16,7 @@
// under the License.
suite("test_mask_function") {
+ sql " set enable_fallback_to_original_planner = false; "
sql """
drop table if exists table_mask_test;
"""
@@ -75,4 +76,21 @@ suite("test_mask_function") {
qt_select_digital_masking """
select digital_masking(13812345678);
"""
+
+ test {
+ sql """ select mask_last_n("12345", -100); """
+ exception "function mask_last_n only accept non-negative input for 2nd
argument but got -100"
+ }
+ test {
+ sql """ select mask_first_n("12345", -100); """
+ exception "function mask_first_n only accept non-negative input for
2nd argument but got -100"
+ }
+ test {
+ sql """ select mask_last_n("12345", id) from table_mask_test; """
+ exception "mask_last_n must accept literal for 2nd argument"
+ }
+ test {
+ sql """ select mask_first_n("12345", id) from table_mask_test; """
+ exception "mask_first_n must accept literal for 2nd argument"
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org
For additional commands, e-mail: commits-h...@doris.apache.org
