korlov42 commented on code in PR #3335:
URL: https://github.com/apache/ignite-3/pull/3335#discussion_r1512718731
##########
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;
Review Comment:
why do we need all three fields? Is it not enough to collect everything in a
single `List<List<RexNode>>`?
##########
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) {
Review Comment:
why DELETE is different from UPDATE? Btw, what about MERGE? Do we have
pruning tests for MERGE?
##########
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:
this doesn't make any sense. Do you have tests for INSERT with rearranged
column list?
##########
modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningMetadataTest.java:
##########
@@ -100,11 +101,76 @@ public class PartitionPruningMetadataTest extends
AbstractPlannerTest {
.distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
.build());
- /** Basic test cases for partition pruning metadata extractor. */
- @ParameterizedTest
+ /** Basic test cases for partition pruning metadata extractor, select
case. */
+ @ParameterizedTest(name = "SELECT: {0}")
@EnumSource(TestCaseBasic.class)
- public void testBasic(TestCaseBasic testCaseSimple) {
- checkPruningMetadata(testCaseSimple.data);
+ public void testBasicSelect(TestCaseBasic testCaseSimple) {
+ checkPruningMetadata(testCaseSimple.data, SqlKind.SELECT);
+ }
+
+ /** Basic test cases for partition pruning metadata extractor, delete
case. */
+ @ParameterizedTest(name = "DELETE: {0}")
+ @EnumSource(TestCaseBasic.class)
+ public void testBasicDelete(TestCaseBasic testCaseSimple) {
+ checkPruningMetadata(testCaseSimple.data, SqlKind.DELETE);
+ }
+
+ /** Basic test cases for partition pruning metadata extractor, delete
case. */
+ @ParameterizedTest(name = "INSERT: {0}")
+ @EnumSource(TestCaseBasicInsert.class)
+ public void testBasicInsert(TestCaseBasicInsert testCaseSimple) {
+ checkPruningMetadata(testCaseSimple.data, SqlKind.INSERT);
+ }
+
+ /** Basic test cases for partition pruning metadata extractor, delete
case. */
+ @ParameterizedTest(name = "UPDATE: {0}")
+ @EnumSource(TestCaseBasicUpdate.class)
+ public void testBasicInsert(TestCaseBasicUpdate testCaseSimple) {
Review Comment:
we have to decided whether this test covers DELETE, UPDATE, or INSERT case
##########
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)) {
+ projList.add(((RexInputRef) node).getIndex());
+ }
+ } else {
+ if (!isValueExpr(node)) {
+ return;
+ }
+ }
+ }
+
+ if (!projList.containsAll(keysList)) {
+ return;
+ }
+
+ projections = replaceInputRefs(projections);
+
+ mapping = RexUtils.inversePermutation(projections,
+ table.getRowType(Commons.typeFactory()), true);
+ }
+
+ List<RexNode> andEqNodes = new ArrayList<>();
+
+ RelDataType rowTypes = table.getRowType(Commons.typeFactory());
+
+ List<List<RexNode>> dataRows = vals != null ?
Commons.cast(vals.getTuples()) : exprProjections;
+
+ for (List<RexNode> items : dataRows) {
+ List<RexNode> andNodes = new ArrayList<>(keysList.size());
+ for (int key : keysList) {
+ RexLocalRef ref;
+ if (projections != null) {
+ ref = (RexLocalRef)
projections.get(mapping.getSourceOpt(key));
+ } else {
+ ref =
rexBuilder.makeLocalRef(rowTypes.getFieldList().get(key).getType(), key);
+ }
+ RexNode lit = items.get(mapping != null ?
mapping.getSourceOpt(key) : key);
+ if (!isValueExpr(lit)) {
+ return;
+ }
+ RexNode eq = rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
ref, lit);
+ andNodes.add(eq);
+ }
+ if (andNodes.size() > 1) {
+ assert andNodes.size() % 2 == 0;
+ RexNode node0 = rexBuilder.makeCall(SqlStdOperatorTable.AND,
andNodes);
+ andEqNodes.add(node0);
+ } else {
+ andEqNodes.add(andNodes.get(0));
+ }
+ }
+
+ if (!nullOrEmpty(andEqNodes)) {
+ RexNode call = rexBuilder.makeCall(SqlStdOperatorTable.OR,
andEqNodes);
Review Comment:
in case of single-tuple insert exception is thrown here
```
java.lang.AssertionError: wrong operand count 1 for OR
at org.apache.calcite.util.Litmus.lambda$static$0(Litmus.java:31)
at
org.apache.calcite.sql.SqlBinaryOperator.validRexOperands(SqlBinaryOperator.java:227)
at org.apache.calcite.rex.RexCall.<init>(RexCall.java:85)
at org.apache.calcite.rex.RexBuilder.makeCall(RexBuilder.java:272)
at
org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningMetadataExtractor.extractFromValues(PartitionPruningMetadataExtractor.java:290)
```
we should add such tests
##########
modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningPredicate.java:
##########
@@ -89,6 +89,17 @@ public ColocationGroup prunePartitions(ColocationGroup
colocationGroup) {
Map<String, List<PartitionWithConsistencyToken>> partitionsPerNode =
newHashMap(colocationGroup.nodeNames().size());
Set<String> newNodes = new HashSet<>();
+ List<NodeWithConsistencyToken> newAssignments = new
ArrayList<>(colocationGroup.assignments().size());
+
+ for (int p = 0; p < colocationGroup.assignments().size(); p++) {
+ NodeWithConsistencyToken nodeWithConsistencyToken =
colocationGroup.assignments().get(p);
+
+ if (remainingPartitions.contains(p)) {
+ newAssignments.add(nodeWithConsistencyToken);
+ } else {
+ newAssignments.add(null);
Review Comment:
this doesn't look safe. Let's replace list with mapping partition id to
NodeWithConsistencyToken
--
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]
