Extend native type rewrite for cast, function call, and simple case expression #25142
Conversation
…ning custom types
presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java
Outdated
Show resolved
Hide resolved
| ? new FunctionCall(node.getLocation().get(), QualifiedName.of(FUNCTION_ELEMENT_AT), node.getWindow(), node.getFilter(), node.getOrderBy(), node.isDistinct(), node.isIgnoreNulls(), arguments) | ||
| : new FunctionCall(QualifiedName.of(FUNCTION_ELEMENT_AT), node.getWindow(), node.getFilter(), node.getOrderBy(), node.isDistinct(), node.isIgnoreNulls(), arguments); | ||
| else { | ||
| arguments.set(0, rewriteIfCastOrDereferenceExpression(argument)); |
There was a problem hiding this comment.
arguments will be immutable. Need to make a new list.
| // Return node without rewriting. | ||
| return node; | ||
| } | ||
| argument = rewriteIfCastOrDereferenceExpression(argument); |
There was a problem hiding this comment.
instead of this custom rewriteIfCastOrdereferenceExpression, why don't we visit all the argument s
| argumentType = ((TypeWithName) argumentType).getType(); | ||
| QualifiedName functionName = node.getName(); | ||
| List<Expression> arguments = node.getArguments(); | ||
| if (node.getArguments().size() == 1) { |
There was a problem hiding this comment.
why do we specifically check for 1 argument? why not visit all the arguments and rewrite as needed?
| argumentType = functionAndTypeResolver.getType(parseTypeSignature(((DereferenceExpression) argument).getBase().toString())); | ||
| } | ||
| else { | ||
| // ENUM_KEY is only supported with Cast or DereferenceExpression for now. |
There was a problem hiding this comment.
I don't know if there is a reason why, but based on all the cases I've encountered, custom types were either a cast expression or dereference expression
There was a problem hiding this comment.
That's most common, but not sure it's guaranteed. The rewriter should be able to support all query shapes. If we leave an enum_key function in, queries will fail during execution.
| return peelIfDereferenceExpression(argument); | ||
| } | ||
|
|
||
| private Expression peelIfDereferenceExpression(Expression argument) |
There was a problem hiding this comment.
Are there any downsides to unwrapping the dereference expression? Can we just have this happen always? i.e. have a visitDereferenceExpression() that handles it. (or have visitExpression call to an ExpressionTreeRewriter that handles all the expression rewrites including Dereference, Cast, SimpleCaseExpression, etc. that way we'll automatically recurse into e.g. cast contained in an argument to some other custom function no matter how many layers deep.)
| } | ||
| } | ||
| } | ||
| catch (IllegalArgumentException | UnknownTypeException e) { |
There was a problem hiding this comment.
in what cases does the rewrite fail. Could we write the code so it doesn't happen instead of relying on exception handling, which could mask bugs.
There was a problem hiding this comment.
I found that it can fail while getting the type (functionAndTypeResolver.getType) for other dereference expressions, for example for a named field in a row or getting some values from a CTE, as it is not recognized as a valid type. I'm not sure what else we can do if this throws other than return the original node. Would logging help?
| } | ||
|
|
||
| @Override | ||
| protected Node visitSimpleCaseExpression(SimpleCaseExpression node, Void context) |
There was a problem hiding this comment.
i think if you use an ExpressionTreeRewriter for rewriting expressions generally, you won't need to add custom logic for handling this or other expressions that might have arguments that have enum functions. And it will work better because it can be nested as many layers deep as someone cares to write.
|
Noob question, why do we rewrite the expression tree rather than the plan fragment? |
agree. but it won't be trivial to set up a mock type manager which provides custom types. |
|
@rschlussel Thanks for the review! I think by calling visitDereferenceExpression or visitExpression and letting that handle the resolving of the custom types will address most of your comments, will update the PR with it |
I think either way works? |
This PR extends #24916 by adding support for additional cases in the NativeExecutionTypeRewrite class. Assuming enum types appear as either Cast or DereferenceExpression expressions within NativeExecutionTypeRewrite, implemented logic for rewriting these expressions into their base literal expressions.
Specifically the updates from #24916 are: