zstan commented on code in PR #3335:
URL: https://github.com/apache/ignite-3/pull/3335#discussion_r1514330927
##########
modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadataExtractor.java:
##########
@@ -137,6 +150,153 @@ public IgniteRel visit(IgniteTableScan rel) {
return rel;
}
+ private static class ModifyNodeShuttle extends IgniteRelShuttle {
+ List<RexNode> projections = null;
+ IgniteValues values = null;
+ List<List<RexNode>> exprProjections = null;
+ boolean collectExprProjections = false;
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteProject rel) {
+ if (projections == null && !collectExprProjections) {
+ projections = rel.getProjects();
+ } else {
+ if (exprProjections == null) {
+ exprProjections = new ArrayList<>();
+ }
+ exprProjections.add(rel.getProjects());
+ }
+ return super.visit(rel);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteValues rel) {
+ if (!collectExprProjections) {
+ values = rel;
+ }
+ return super.visit(rel);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteUnionAll rel) {
+ collectExprProjections = true;
+ return super.visit(rel);
+ }
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteTableModify rel) {
+ if (rel.getOperation() != INSERT && rel.getOperation() != DELETE) {
+ return super.visit(rel);
+ }
+
+ IgniteTable table = rel.getTable().unwrap(IgniteTable.class);
+
+ RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
+
+ ModifyNodeShuttle modify = new ModifyNodeShuttle();
+ rel.accept(modify);
+
+ extractFromValues(rel.sourceId(), table, modify.values,
modify.projections, modify.exprProjections, rexBuilder);
+
+ return super.visit(rel);
+ }
+
+ private void extractFromValues(
+ long sourceId,
+ IgniteTable table,
+ @Nullable IgniteValues vals,
+ @Nullable List<RexNode> projections,
+ @Nullable List<List<RexNode>> exprProjections,
+ RexBuilder rexBuilder
+ ) {
+ if (vals == null && exprProjections == null) {
+ return;
+ }
+
+ IgniteDistribution distribution = table.distribution();
+ if (!distribution.function().affinity()) {
+ return;
+ }
+
+ IntArrayList keysList = new
IntArrayList(distribution.getKeys().size());
+ for (Integer key : distribution.getKeys()) {
+ keysList.add(key.intValue());
+ }
+
+ Mappings.TargetMapping mapping = null;
+
+ if (projections != null) {
+ IntArrayList projList = new IntArrayList();
+ for (RexNode node : projections) {
+ if (node instanceof RexInputRef) {
+ int idx = ((RexInputRef) node).getIndex();
+ if (keysList.contains(idx)) {
Review Comment:
yes, check PartitionPruningMetadataTest.TestCaseBasicInsert#SIMPLE_1j1
--
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]
