Skip to content

branch-4.2: [fix](ai) Correct resource usage #70297

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
linrrzqqq wants to merge 1 commit into
base: branch-4.2
Choose a base branch
from
Open

branch-4.2: [fix](ai) Correct resource usage #70297

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,9 +77,11 @@
import org.apache.doris.nereids.trees.expressions.functions.FunctionBuilder;
import org.apache.doris.nereids.trees.expressions.functions.RewriteWhenAnalyze;
import org.apache.doris.nereids.trees.expressions.functions.Udf;
import org.apache.doris.nereids.trees.expressions.functions.agg.AIAgg;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.agg.NullableAggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.agg.SupportMultiDistinct;
import org.apache.doris.nereids.trees.expressions.functions.ai.AIFunction;
import org.apache.doris.nereids.trees.expressions.functions.scalar.ElementAt;
import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda;
import org.apache.doris.nereids.trees.expressions.functions.udf.AliasUdfBuilder;
Expand Down Expand Up @@ -616,13 +618,8 @@ public Expression visitUnboundFunction(UnboundFunction unboundFunction, Expressi

if (wantToParseSqlFromSqlCache) {
sqlCacheContext = context.cascadesContext.getStatementContext().getSqlCacheContext();
if (builder instanceof AliasUdfBuilder
|| buildResult.second instanceof Udf) {
if (sqlCacheContext.isPresent()) {
sqlCacheContext.get().setCannotProcessExpression(true);
}
}
}
disableSqlCacheForUnsupportedFunction(buildResult.second);
if (unboundFunction.isSkew() && unboundFunction.isDistinct() && buildResult.first instanceof AggregateFunction
&& buildResult.first instanceof SupportMultiDistinct) {
Expression bound = ((AggregateFunction) buildResult.first).withIsSkew(true);
Expand All @@ -647,9 +644,18 @@ public Expression visitUnboundFunction(UnboundFunction unboundFunction, Expressi
@Override
public Expression visitBoundFunction(BoundFunction boundFunction, ExpressionRewriteContext context) {
boundFunction = (BoundFunction) super.visitBoundFunction(boundFunction, context);
disableSqlCacheForUnsupportedFunction(boundFunction);
return TypeCoercionUtils.processBoundFunction(boundFunction);
}

private void disableSqlCacheForUnsupportedFunction(BoundFunction function) {
if (wantToParseSqlFromSqlCache
&& (function instanceof Udf || function instanceof AIFunction || function instanceof AIAgg)) {
getCascadesContext().getStatementContext().getSqlCacheContext()
.ifPresent(cache -> cache.setCannotProcessExpression(true));
}
}

@Override
public Expression visitWindow(WindowExpression windowExpression, ExpressionRewriteContext context) {
windowExpression = (WindowExpression) super.visitWindow(windowExpression, context);
Expand Down
7 changes: 7 additions & 0 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
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.ai.AIFunction;
import org.apache.doris.nereids.trees.expressions.literal.StringLiteral;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.types.StringType;
Expand Down Expand Up @@ -73,6 +74,11 @@ private AIAgg(NullableAggregateFunctionParams functionParams) {
super(functionParams);
}

@Override
public boolean isDeterministic() {
return false;
}

@Override
public void checkLegalityAfterRewrite() {
// The rewrite has folded a constant task and resource name, unless constant folding is skipped
Expand Down Expand Up @@ -113,6 +119,7 @@ public void checkLegalityBeforeTypeCoercion() {
private void checkResource(Expression resourceNameArgument) {
//Check if the resource is valid
String resourceName = resourceNameArgument.toString().replaceAll("^['\"]|['\"]$", "");
AIFunction.checkResourceUsagePrivilege(resourceName);
Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName);
if (!(resource instanceof AIResource)) {
throw new AnalysisException("AI resource '" + resourceName + "' does not exist");
Expand Down
16 changes: 16 additions & 0 deletions ...ore/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import org.apache.doris.catalog.AIResource;
import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.Resource;
import org.apache.doris.common.ErrorCode;
import org.apache.doris.mysql.privilege.PrivPredicate;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
Expand Down Expand Up @@ -48,6 +50,11 @@ public AIFunction(String name, Expression... expressions) {

public abstract int getMaxArgsNum();

@Override
public boolean isDeterministic() {
return false;
}

@Override
public void checkLegalityAfterRewrite() {
// The rewrite has folded a constant resource name, unless constant folding is skipped
Expand All @@ -66,6 +73,7 @@ private void checkResource(Expression resourceNameArgument) {

//Check if the resource is valid
String resourceName = resourceNameArgument.toString().replaceAll("^['\"]|['\"]$", "");
checkResourceUsagePrivilege(resourceName);
Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName);
if (!(resource instanceof AIResource)) {
throw new AnalysisException("AI resource '" + resourceName + "' does not exist");
Expand All @@ -77,6 +85,14 @@ private void checkResource(Expression resourceNameArgument) {
Resource.registerUsedAIResourceName(resourceName);
}

public static void checkResourceUsagePrivilege(String resourceName) {
if (!Env.getCurrentEnv().getAccessManager()
.checkResourcePriv(ConnectContext.get(), resourceName, PrivPredicate.USAGE)) {
throw new AnalysisException(ErrorCode.ERR_RESOURCE_ACCESS_DENIED_ERROR.formatErrorMsg(
PrivPredicate.USAGE.getPrivs().toString(), resourceName));
}
}

/**
* If a user doesn't specify which resource to use when calling a function
* the current session variable will automatically allocate a resource.
Expand Down
1 change: 1 addition & 0 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ private static String requireStringLiteral(Expression arg, String argName, Strin
}

private static void validateAIResource(String resourceName, boolean isMultimodal) {
checkResourceUsagePrivilege(resourceName);
Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName);
if (!(resource instanceof AIResource)) {
throw new AnalysisException("AI resource '" + resourceName + "' does not exist");
Expand Down
Loading
Loading