[vectorized](function) support array_first_index function (#18175)

mysql> select array_first_index(x->x+1>3, [2, 3, 4]);
+-------------------------------------------------------------------+
| array_first_index(array_map([x] -> x(0) + 1 > 3, ARRAY(2, 3, 4))) |
+-------------------------------------------------------------------+
|                                                                 2 |
+-------------------------------------------------------------------+

mysql> select array_first_index(x -> x is null, [null, 1, 2]);
+----------------------------------------------------------------------+
| array_first_index(array_map([x] -> x(0) IS NULL, ARRAY(NULL, 1, 2))) |
+----------------------------------------------------------------------+
|                                                                    1 |
+----------------------------------------------------------------------+

mysql> select array_first_index(x->power(x,2)>10, [1, 2, 3, 4]);
+---------------------------------------------------------------------------------+
| array_first_index(array_map([x] -> power(x(0), 2.0) > 10.0, ARRAY(1, 2, 3, 4))) |
+---------------------------------------------------------------------------------+
|                                                                               4 |
+---------------------------------------------------------------------------------+
This commit is contained in:
gitccl
2023-03-31 12:51:29 +08:00
committed by GitHub
parent 307170030c
commit 20b3bdb000
10 changed files with 427 additions and 4 deletions

View File

@ -35,12 +35,12 @@ import java.util.List;
public class LambdaFunctionCallExpr extends FunctionCallExpr {
public static final ImmutableSet<String> LAMBDA_FUNCTION_SET = new ImmutableSortedSet.Builder(
String.CASE_INSENSITIVE_ORDER).add("array_map").add("array_filter").add("array_exists").add("array_sortby")
.build();
.add("array_first_index").build();
// The functions in this set are all normal array functions when implemented initially.
// and then wants add lambda expr as the input param, so we rewrite it to contains an array_map lambda function
// rather than reimplementing a lambda function, this will be reused the implementation of normal array function
public static final ImmutableSet<String> LAMBDA_MAPPED_FUNCTION_SET = new ImmutableSortedSet.Builder(
String.CASE_INSENSITIVE_ORDER).add("array_exists").add("array_sortby").build();
String.CASE_INSENSITIVE_ORDER).add("array_exists").add("array_sortby").add("array_first_index").build();
private static final Logger LOG = LogManager.getLogger(LambdaFunctionCallExpr.class);
@ -105,7 +105,8 @@ public class LambdaFunctionCallExpr extends FunctionCallExpr {
throw new AnalysisException(getFunctionNotFoundError(collectChildReturnTypes()));
}
fn.setReturnType(ArrayType.create(lambda.getChild(0).getType(), true));
} else if (fnName.getFunction().equalsIgnoreCase("array_exists")) {
} else if (fnName.getFunction().equalsIgnoreCase("array_exists")
|| fnName.getFunction().equalsIgnoreCase("array_first_index")) {
if (fnParams.exprs() == null || fnParams.exprs().size() < 1) {
throw new AnalysisException("The " + fnName.getFunction() + " function must have at least one param");
}
@ -138,7 +139,6 @@ public class LambdaFunctionCallExpr extends FunctionCallExpr {
LOG.warn("fn {} not exists", this.toSqlImpl());
throw new AnalysisException(getFunctionNotFoundError(collectChildReturnTypes()));
}
fn.setReturnType(getChild(0).getType());
} else if (fnName.getFunction().equalsIgnoreCase("array_filter")) {
if (fnParams.exprs() == null || fnParams.exprs().size() != 2) {
throw new AnalysisException("The " + fnName.getFunction() + " function must have two params");