lowka commented on code in PR #3335:
URL: https://github.com/apache/ignite-3/pull/3335#discussion_r1524483207
##########
modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadataExtractor.java:
##########
@@ -138,6 +152,177 @@ public IgniteRel visit(IgniteTableScan rel) {
return rel;
}
+ private static class ModifyNodeShuttle extends IgniteRelShuttle {
+ List<RexNode> projections = null;
+ List<List<RexNode>> expressions = null;
+ boolean unionRaised = false;
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteProject rel) {
+ if (!unionRaised) {
+ // projection before union
+ assert projections == null;
+ // unexpected branch
+ if (projections != null) {
+ throw Util.FoundOne.NULL;
+ }
+ projections = rel.getProjects();
+ } else {
+ // projections after union
+ if (expressions == null) {
+ expressions = new ArrayList<>();
+ }
+ expressions.add(rel.getProjects());
+ }
+
+ return super.visit(rel);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteValues rel) {
+ if (expressions == null) {
+ expressions = new ArrayList<>();
+ }
+
+ List<List<RexNode>> values = Commons.cast(rel.getTuples());
+
+ expressions.addAll(values);
+
+ return super.visit(rel);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteUnionAll rel) {
+ unionRaised = true;
+ return super.visit(rel);
+ }
+
+ @Override
+ protected void visitChild(IgniteRel parent, int i, IgniteRel child) {
+ if (child instanceof IgniteProject
+ || child instanceof IgniteValues
+ || child instanceof IgniteUnionAll
+ || child instanceof IgniteExchange
+ || child instanceof IgniteTrimExchange) {
+ super.visitChild(parent, i, child);
+ } else {
+ throw Util.FoundOne.NULL;
+ }
+ }
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteTableModify rel) {
+ if (rel.getOperation() != INSERT) {
+ return super.visit(rel);
+ }
+
+ IgniteTable table = rel.getTable().unwrap(IgniteTable.class);
+
+ RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
+
+ ModifyNodeShuttle modify = new ModifyNodeShuttle();
+
+ try {
+ rel.accept(modify);
+ } catch (Util.FoundOne e) {
+ return rel;
+ }
+
+ extractFromValues(rel.sourceId(), table, modify.projections,
modify.expressions, rexBuilder);
+
+ return super.visit(rel);
+ }
+
+ private void extractFromValues(
+ long sourceId,
+ IgniteTable table,
+ @Nullable List<RexNode> projections,
+ @Nullable List<List<RexNode>> expressions,
+ RexBuilder rexBuilder
+ ) {
+ if (expressions == null) {
Review Comment:
this 3 conditions are copied from `extractFromTable`
I think it better to extract these conditions into a single method that
contains all them conditions and make such method to return a `IntKeyList` of
`distribution.getKeys()` columns.
If that list is empty, then there is nothing to do.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
