IGNITE-24719 SQL Calcite: Fix mapping sending for trim exchange - Fixes #11912.

Signed-off-by: Aleksey Plekhanov <plehanov.alex@gmail.com>

Author: alex-plekhanov

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/metadata/FragmentMapping.java

@@ -20,6 +20,7 @@
 import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.UUID;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -167,6 +168,16 @@ else if (grps.size() > 1)
         return F.first(grps);
     }
 
+    /** Create fragment mapping with explicit mapping for groups by source ids. */
+    public FragmentMapping explicitMapping(Set<Long> srcIds) {
+        Set<ColocationGroup> explicitMappingGrps = U.newIdentityHashSet();
+
+        srcIds.forEach(srcId -> explicitMappingGrps.add(findGroup(srcId)));
+
+        return new FragmentMapping(Commons.transform(colocationGroups,
+            g -> explicitMappingGrps.contains(g) ? g.explicitMapping() : g));
+    }
+
     /** {@inheritDoc} */
     @Override public void prepareMarshal(GridCacheSharedContext<?, ?> ctx) {
         colocationGroups.forEach(g -> g.prepareMarshal(ctx));

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/prepare/Fragment.java

@@ -17,7 +17,9 @@
 
 package org.apache.ignite.internal.processors.query.calcite.prepare;
 
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Supplier;
@@ -34,6 +36,7 @@
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteReceiver;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRel;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteSender;
+import org.apache.ignite.internal.processors.query.calcite.rel.IgniteTrimExchange;
 import org.apache.ignite.internal.processors.query.calcite.trait.IgniteDistributions;
 import org.apache.ignite.internal.util.tostring.GridToStringExclude;
 import org.apache.ignite.internal.util.typedef.internal.S;
@@ -153,6 +156,22 @@ private FragmentMapping mapping(MappingQueryContext ctx, RelMetadataQuery mq, Su
                 else
                     mapping = FragmentMapping.create(ctx.localNodeId()).colocate(mapping);
             }
+            else {
+                // For remote fragments we should ensure that mapping is explicitely sent to remote nodes if
+                // it's required by trim exchange.
+                Set<Long> trimSrcIds = new HashSet<>();
+
+                root.accept(new IgniteRelShuttle() {
+                    @Override public IgniteRel visit(IgniteTrimExchange rel) {
+                        trimSrcIds.add(rel.sourceId());
+
+                        return super.visit(rel);
+                    }
+                });
+
+                if (!trimSrcIds.isEmpty())
+                    mapping = mapping.explicitMapping(trimSrcIds);
+            }
 
             if (single() && mapping.nodeIds().size() > 1) {
                 // this is possible when the fragment contains scan of a replicated cache, which brings

…tegration/JoinRehashIntegrationTest.java …tion/DistributedJoinIntegrationTest.javamodules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/JoinRehashIntegrationTest.java renamed to modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/DistributedJoinIntegrationTest.java modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/JoinRehashIntegrationTest.java renamed to modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/DistributedJoinIntegrationTest.java

@@ -23,15 +23,15 @@
 import org.junit.Test;
 
 /** */
-public class JoinRehashIntegrationTest extends AbstractBasicIntegrationTransactionalTest {
+public class DistributedJoinIntegrationTest extends AbstractBasicIntegrationTransactionalTest {
     /** {@inheritDoc} */
     @Override protected int nodeCount() {
         return 3;
     }
 
     /** Test that resources (in particular inboxes) are cleaned up after executing join with rehashing. */
     @Test
-    public void testResourceCleanup() throws Exception {
+    public void testRehashResourceCleanup() throws Exception {
         prepareTables();
 
         String sql = "SELECT sum(i.price * i.amount)" +
@@ -55,7 +55,7 @@ public void testResourceCleanup() throws Exception {
 
     /** Tests that null values are filtered out on rehashing. */
     @Test
-    public void testNullAffinityKeys() {
+    public void testRehashNullAffinityKeys() {
         prepareTables();
 
         // Add null values.
@@ -86,6 +86,20 @@ public void testRehashOnRightHand() {
             .check();
     }
 
+    /** */
+    @Test
+    public void testTrimExchange() {
+        sql("CREATE TABLE order_ids(id INTEGER PRIMARY KEY) WITH TEMPLATE=REPLICATED," + atomicity());
+        prepareTables();
+
+        sql("INSERT INTO order_ids(id) SELECT id FROM orders");
+
+        assertQuery("SELECT sum(o.id) FROM orders o JOIN order_ids oid ON (o.id = oid.id)")
+            .matches(QueryChecker.containsSubPlan("IgniteTrimExchange"))
+            .returns(435L)
+            .check();
+    }
+
     /** Prepare tables orders and order_items with data. */
     private void prepareTables() {
         sql("CREATE TABLE items (\n" +

modules/calcite/src/test/java/org/apache/ignite/testsuites/IntegrationTestSuite.java

@@ -30,6 +30,7 @@
 import org.apache.ignite.internal.processors.query.calcite.integration.CorrelatesIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.DataTypesTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.DateTimeTest;
+import org.apache.ignite.internal.processors.query.calcite.integration.DistributedJoinIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.DynamicParametersIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.ExpiredEntriesIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.FunctionsTest;
@@ -42,7 +43,6 @@
 import org.apache.ignite.internal.processors.query.calcite.integration.IndexSpoolIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.IntervalTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.JoinIntegrationTest;
-import org.apache.ignite.internal.processors.query.calcite.integration.JoinRehashIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.KeepBinaryIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.KillCommandDdlIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.KillQueryCommandDdlIntegrationTest;
@@ -145,7 +145,7 @@
     ExpiredEntriesIntegrationTest.class,
     TimeoutIntegrationTest.class,
     PartitionPruneTest.class,
-    JoinRehashIntegrationTest.class,
+    DistributedJoinIntegrationTest.class,
     IndexWithSameNameCalciteTest.class,
     AuthorizationIntegrationTest.class,
     DdlTransactionCalciteSelfTest.class,