Support pushing dereferences within lambdas into table scan #23148
Conversation
|
This pull request has gone a while without any activity. Tagging the Trino developer relations team: @bitsondatadev @colebow @mosabua |
| @@ -122,6 +122,8 @@ public class FeaturesConfig | |||
|
|
|||
| private boolean faultTolerantExecutionExchangeEncryptionEnabled = true; | |||
|
|
|||
| private boolean pushFieldDereferenceLambdaIntoScanEnabled; | |||
There was a problem hiding this comment.
this should be with other optimizer properties in OptimizerConfig
TBH: I think it should be connector property. We already have io.trino.plugin.iceberg.IcebergConfig#projectionPushdownEnabled so it might not be needed at all
| return pushFieldDereferenceLambdaIntoScanEnabled; | ||
| } | ||
|
|
||
| @Config("experimental.enable-push-field-dereference-lambda-into-scan.enabled") |
There was a problem hiding this comment.
I don't think it should really be experimental.
TBH I think it should be connector property
| @@ -694,6 +706,37 @@ else if (functionName.equals(builtinFunctionName(MODULUS))) { | |||
| new io.trino.spi.expression.Call(node.type(), MODULUS_FUNCTION_NAME, ImmutableList.of(left, right)))); | |||
| } | |||
|
|
|||
| // Very narrow case that only tries to extract a particular type of lambda expression | |||
| // TODO: Expand the scope | |||
There was a problem hiding this comment.
Please create Trino issue which would track the TODO and reference it here.
| // Very narrow case that only tries to extract a particular type of lambda expression | ||
| // TODO: Expand the scope | ||
| if (translateArrayFieldReference && functionName.equals(builtinFunctionName(ARRAY_TRANSFORM_NAME))) { | ||
| List<Expression> allNodeArgument = node.arguments(); |
| List<Expression> allNodeArgument = node.arguments(); | ||
| // at this point, SubscriptExpression should already been pushed down by PushProjectionIntoTableScan, | ||
| // if not, it means its referenced by other expressions. we only care about SymbolReference at this moment | ||
| List<Expression> inputExpressions = allNodeArgument.stream().filter(Reference.class::isInstance) |
There was a problem hiding this comment.
Please check arguments.getFirst() explicitly, see SpecializeTransformWithJsonParse
| if (lambdaExpressions.get(0).body() instanceof Row row) { | ||
| List<Expression> rowFields = row.items(); | ||
| List<ConnectorExpression> translatedRowFields = | ||
| rowFields.stream().map(e -> process(e)).filter(Optional::isPresent).map(Optional::get).collect(toImmutableList()); |
There was a problem hiding this comment.
nit:
List<ConnectorExpression> translatedRowFields = rowFields.stream()
.map(this::process)
.filter(Optional::isPresent)
.collect(toImmutableList());
|
|
||
| // This class is used to represent expression with dereferences into Array | ||
| // Target is the actual reference to the array. elementFieldDereferences are the field dereferences | ||
| public class ArrayFieldDereference |
There was a problem hiding this comment.
This is a bit confusing for me.
Let's assume that the table schema is:
field1 <- bigint
field2 <- array(row(x:bigint y:bigint, z:bigint))
What does
ArrayFieldDereference(field2, [x]) produce?
Is it array(row(x)) or array(x)?
what does
ArrayFieldDereference(field2, [x, y]) produce?
Is it array(row(x, y))?
| public ArrayFieldDereference(Type type, ConnectorExpression target, List<ConnectorExpression> elementFieldDereference) | ||
| { | ||
| super(type); | ||
| checkArgument(type instanceof ArrayType, "wrong input type for ArrayFieldDereference"); |
There was a problem hiding this comment.
Please add a check:
checkArgument(!elementFieldDereference.isEmpty(), ...)
checkArgument(elementFieldDereference.size() == 1 || ((ArrayType) type).getElementType() instanceOf Row, ...)
| @@ -999,6 +1005,15 @@ public PlanOptimizers( | |||
| new PushPartialAggregationThroughExchange(plannerContext), | |||
| new PruneJoinColumns(), | |||
| new PruneJoinChildrenColumns()))); | |||
| // This rule does not touch query plans, but only add subfields if necessary. Trigger at the near end | |||
There was a problem hiding this comment.
Are there queries that actually fail to get pushdown if we don't have this rule here? Maybe it's not needed?
| @@ -132,4 +143,113 @@ private static boolean prefixExists(Expression expression, Set<Expression> expre | |||
| verify(current instanceof Reference); | |||
| return false; | |||
| } | |||
|
|
|||
There was a problem hiding this comment.
nit: public methods above private methods
| */ | ||
| public static Optional<Reference> getSubscriptLambdaInputExpression(Expression expression) | ||
| { | ||
| if (expression instanceof Call functionCall) { |
There was a problem hiding this comment.
nit: if you revert it:
if (!(expression instanceOf Call functionCall) ||
!functionCall.equals(builtinFunctionName(ARRAY_TRANSFORM_NAME))) {
return Optional.empty();
}
// rest of code
that the happy path does not have extra indents.
| CatalogSchemaFunctionName functionName = functionCall.function().name(); | ||
|
|
||
| if (functionName.equals(builtinFunctionName(ARRAY_TRANSFORM_NAME))) { | ||
| List<Expression> allNodeArgument = functionCall.arguments(); |
There was a problem hiding this comment.
similar pattern happens in ConnectorExpressionTranslator. Maybe we can extract something like
record ArrayDereferenceExpression(Expression argument, Lambda lambda)
then we could have method Optional<ArrayDereferenceExpression> getArrayDereferenceExpression(Expression expression)
and use it in rules and in ConnectorExpressionTranslator?
| /** | ||
| * Extract the sub-expressions of type {@link Reference} from the {@param expression} | ||
| */ | ||
| public static List<Reference> getReferences(Expression expression) |
There was a problem hiding this comment.
I feel like this you should use io.trino.sql.planner.SymbolsExtractor#extractUnique(io.trino.sql.ir.Expression) instead. This is because lambdas nested within an expression can reference lamba arguments, see: io.trino.sql.planner.SymbolsExtractor.SymbolBuilderVisitor
| /** | ||
| * Extract the sub-expressions of type {@link Reference} and subscript lambda {@link FunctionCall} from the {@param expression} | ||
| */ | ||
| private static Map<Expression, Reference> getSymbolReferencesAndSubscriptLambdas(Expression expression) |
There was a problem hiding this comment.
I don't think it does what it means to. See io.trino.sql.planner.SymbolsExtractor.SymbolBuilderVisitor. Symbols within lambdas need to be handled carefully
| /** | ||
| * Transforms: | ||
| * <pre> | ||
| * Project(c := f(a, x -> x[1]), d := g(b)) |
There was a problem hiding this comment.
I think we should limit it to RowType dereference, e.g:
if (!(node.base().type() instanceof RowType)) {
return Optional.empty();
}
because support for array dereference is generally harder
| import static java.util.Objects.requireNonNull; | ||
| import static java.util.stream.Collectors.joining; | ||
|
|
||
| // This class is used to represent expression with dereferences into Array |
There was a problem hiding this comment.
I think this could generally be simpler, e.g: we don't really need private final List<ConnectorExpression> elementFieldDereferences;. We could just use combination of
io.trino.spi.expression.Call and
singular io.trino.spi.expression.FieldDereference.
Consider example:
foo:=transform(x, row(x.fieldA, x.fieldB))
it can generally be transformed into
foo:=$arrayOfRows(arrayFieldA, arrayFieldB) -> array<row<fieldA, fieldB>>
|
arrayFieldA:=transform(x, x.fieldA) -> array<fieldA>
arrayFieldB:=transform(x, x.fieldB) -> array<fieldB>
then each transform(x, x.fieldA) and transform(x, x.fieldB) could be pushed separately. Engine then would optimized method $arrayOfRows which would call io.trino.spi.block.RowBlock#fromFieldBlocks.
This way connectors would be much less complex and we could keep using
private final List<Integer> path;
from IcebergColumnHandle (and similar from HiveColumnProjectionInfo). Essentially, de-reference path expressed as index array would be sufficient. I would guess we could also avoid SubField and SubfieldTokenizer.
TBH. I don't think we need to support transform(x, row(x.fieldA, x.fieldB)) case initially. We could add support later on
|
This pull request has gone a while without any activity. Tagging the Trino developer relations team: @bitsondatadev @colebow @mosabua |
Description
original PR: #21957
This is to extend the enhancement discussed here #3925, and depends/extends on the original PR #4270 that is currently rebasing by @Desmeister
Since the issue and discussion had been idled for years and this kind of optimization could be critical to anyone having highly nested schema and using Unnest, I would like to use this PR to formally restart the discussion on how the community want to eventually support this and if this is on the right direction (I have a working version locally, not this one, that speeds up the query while reducing actual data processed)
From my understanding of the previous discussions, this should be done through below steps:
dereferencinginvolved withUnnestinto lambda functions with subscript expressions for each of theUnnestsTableScan, in this case, this rule will help to pushdown the dereferencing further, while for any connectors that dont support dereferencing, the rule will preserve the Lambda expression to remove columnsPushDownDereferenceThroughUnnestand many other expression specific rules.PushDownDereferenceThroughUnnestis not handling any unnest symbols currently, but only replicated symbols. In order to support unnest symbols, I believe at least a new expression has to be created, or subscript expression has to be extended otherwise I dont see an easy way to represent the dereferences so it can be further pushed down through other unnests in anyway. I need more guidance on how this could be done or possible with what we have now, that is why this PR in particular is not handling any complex cases like nested Unnest and only push lambdas down through project and filters in a limited way.dereferencingintoTableScanvisitFunctionCallinConnectorExpressionTranslatorto create a new connector expression (can be merged with existingFieldDereferenceexpression if possible), then passing those into existingapplyProjectionmethod to let connectors decide how to handle those. For this PR, onlyHiveMetadatahas implementation to handle those, other connectors will simply ignoring them. TheapplyProjectionwill create new projections andHiveColumnHandlefor Hive with extendedHiveColumnProjectionInfo.Structs.This PR is written in a way to reduce the impacts to the existing features while I can fully validate the performance impact while gathering feedbacks and directions from the community. Therefore implementations are normally wrapped in an
ifinstead of fully refactoring the existing methodI believe if this is the right direction, changes can be contributed through below phases
Array<dereferences>withinHiveColumnProjectionInfotoSubfieldsor anything similar to that and make sure all methods that used to depend onArray<dereferences>now depend on the new representationapplyProjectionmethod (or not? It can simply be a non-iterative visitor at the very end like now.)The change has been fully validated except rebasing to the latest Trino release that could have a lot of conflicts due to AST/IR refactoring
Byte scanned decreased from 423B to 359B for the sample query, we've seen large performance improvement in production queries
Additional context and related issues
I would really appreciate any kind of comments or feedbacks as without clear directions, I can't further extend this without risking of throwing everything away later. Any of the component should be easily plug in if we have a clear idea of how we want to do it otherwise.
Release notes
( ) This is not user-visible or is docs only, and no release notes are required.
( ) Release notes are required. Please propose a release note for me.
(X) Release notes are required, with the following suggested text: