zstan commented on code in PR #3335:
URL: https://github.com/apache/ignite-3/pull/3335#discussion_r1525805750
##########
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:
done
##########
modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/framework/TestBuilders.java:
##########
@@ -127,7 +131,7 @@
/**
* A collection of builders to create test objects.
*/
-public class TestBuilders {
+public class TestBuilders extends BaseIgniteAbstractTest {
Review Comment:
got it
--
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]
