IGNITE-19998 SQL Calcite: Add support of setting partitions in SqlFieldsQuery (apache#10870)

Author: ivandasch

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

@@ -564,6 +564,7 @@ private <T> T processQuery(
             params,
             qryCtx,
             fldsQry != null && fldsQry.isLocal(),
+            fldsQry != null ? fldsQry.getPartitions() : null,
             exchangeSvc,
             (q, ex) -> qryReg.unregister(q.id(), ex),
             log,

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

@@ -107,6 +107,7 @@ public RootQuery(
         Object[] params,
         QueryContext qryCtx,
         boolean isLocal,
+        int[] parts,
         ExchangeService exch,
         BiConsumer<Query<RowT>, Throwable> unregister,
         IgniteLogger log,
@@ -144,6 +145,7 @@ public RootQuery(
                     .build()
             )
             .local(isLocal)
+            .partitions(parts)
             .logger(log)
             .build();
     }
@@ -157,7 +159,7 @@ public RootQuery(
      * @param schema new schema.
      */
     public RootQuery<RowT> childQuery(SchemaPlus schema) {
-        return new RootQuery<>(sql, schema, params, QueryContext.of(cancel), ctx.isLocal(), exch, unregister, log,
+        return new RootQuery<>(sql, schema, params, QueryContext.of(cancel), ctx.isLocal(), ctx.partitions(), exch, unregister, log,
             plannerTimeout, totalTimeout);
     }
 

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/ExecutionServiceImpl.java

@@ -75,8 +75,10 @@
 import org.apache.ignite.internal.processors.query.calcite.message.QueryStartRequest;
 import org.apache.ignite.internal.processors.query.calcite.message.QueryStartResponse;
 import org.apache.ignite.internal.processors.query.calcite.metadata.AffinityService;
+import org.apache.ignite.internal.processors.query.calcite.metadata.ColocationMappingException;
 import org.apache.ignite.internal.processors.query.calcite.metadata.FragmentDescription;
 import org.apache.ignite.internal.processors.query.calcite.metadata.FragmentMapping;
+import org.apache.ignite.internal.processors.query.calcite.metadata.FragmentMappingException;
 import org.apache.ignite.internal.processors.query.calcite.metadata.MappingService;
 import org.apache.ignite.internal.processors.query.calcite.metadata.RemoteException;
 import org.apache.ignite.internal.processors.query.calcite.prepare.BaseQueryContext;
@@ -564,6 +566,17 @@ private ListFieldsQueryCursor<?> mapAndExecutePlan(
 
         List<Fragment> fragments = plan.fragments();
 
+        if (!F.isEmpty(qry.context().partitions())) {
+            fragments = Commons.transform(fragments, f -> {
+                try {
+                    return f.filterByPartitions(qry.context().partitions());
+                }
+                catch (ColocationMappingException e) {
+                    throw new FragmentMappingException("Failed to calculate physical distribution", f, f.root(), e);
+                }
+            });
+        }
+
         // Local execution
         Fragment fragment = F.first(fragments);
 
@@ -576,7 +589,8 @@ private ListFieldsQueryCursor<?> mapAndExecutePlan(
 
             List<UUID> nodes = mapping.nodeIds();
 
-            assert nodes != null && nodes.size() == 1 && F.first(nodes).equals(localNodeId());
+            assert nodes != null && nodes.size() == 1 && F.first(nodes).equals(localNodeId())
+                    : "nodes=" + nodes + ", localNode=" + localNodeId();
         }
 
         long timeout = qry.remainingTime();

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

@@ -45,7 +45,6 @@
 import org.apache.ignite.plugin.extensions.communication.MessageCollectionItemType;
 import org.apache.ignite.plugin.extensions.communication.MessageReader;
 import org.apache.ignite.plugin.extensions.communication.MessageWriter;
-import org.jetbrains.annotations.NotNull;
 
 /** */
 public class ColocationGroup implements MarshalableMessage {
@@ -114,7 +113,13 @@ public List<UUID> nodeIds() {
      * {@link GridDhtPartitionState#OWNING} state, calculated for distributed tables, involved in query execution.
      */
     public List<List<UUID>> assignments() {
-        return assignments == null ? Collections.emptyList() : assignments;
+        if (assignments != null)
+            return assignments;
+
+        if (!F.isEmpty(nodeIds))
+            return nodeIds.stream().map(Collections::singletonList).collect(Collectors.toList());
+
+        return Collections.emptyList();
     }
 
     /**
@@ -208,39 +213,50 @@ public ColocationGroup colocate(ColocationGroup other) throws ColocationMappingE
     }
 
     /** */
-    public ColocationGroup finalaze() {
-        if (assignments == null && nodeIds == null)
+    public ColocationGroup finalizeMapping() {
+        if (assignments == null)
             return this;
 
-        if (assignments != null) {
+        List<List<UUID>> assignments = new ArrayList<>(this.assignments.size());
+        Set<UUID> nodes = new HashSet<>();
+
+        for (List<UUID> assignment : this.assignments) {
+            UUID first = F.first(assignment);
+            if (first != null)
+                nodes.add(first);
+            assignments.add(first != null ? Collections.singletonList(first) : Collections.emptyList());
+        }
+
+        return new ColocationGroup(sourceIds, new ArrayList<>(nodes), assignments);
+    }
+
+    /** */
+    public ColocationGroup filterByPartitions(int[] parts) {
+        if (!F.isEmpty(assignments)) {
             List<List<UUID>> assignments = new ArrayList<>(this.assignments.size());
             Set<UUID> nodes = new HashSet<>();
-            for (List<UUID> assignment : this.assignments) {
-                UUID first = F.first(assignment);
+
+            if (F.isEmpty(parts))
+                return this;
+
+            for (int i = 0; i < this.assignments.size(); ++i) {
+                UUID first = Arrays.binarySearch(parts, i) >= 0 ? F.first(this.assignments.get(i)) : null;
+
                 if (first != null)
                     nodes.add(first);
-                assignments.add(first != null ? Collections.singletonList(first) : Collections.emptyList());
+
+                assignments.add(first != null ? this.assignments.get(i) : Collections.emptyList());
             }
 
             return new ColocationGroup(sourceIds, new ArrayList<>(nodes), assignments);
         }
 
-        return forNodes0(nodeIds);
+        return this;
     }
 
     /** */
     public ColocationGroup mapToNodes(List<UUID> nodeIds) {
-        return !F.isEmpty(this.nodeIds) ? this : forNodes0(nodeIds);
-    }
-
-    /** */
-    @NotNull private ColocationGroup forNodes0(List<UUID> nodeIds) {
-        List<List<UUID>> assignments = new ArrayList<>(nodeIds.size());
-
-        for (UUID nodeId : nodeIds)
-            assignments.add(Collections.singletonList(nodeId));
-
-        return new ColocationGroup(sourceIds, nodeIds, assignments);
+        return !F.isEmpty(this.nodeIds) ? this : new ColocationGroup(sourceIds, nodeIds, null);
     }
 
     /**
@@ -250,6 +266,9 @@ public ColocationGroup mapToNodes(List<UUID> nodeIds) {
      * @return List of partitions to scan on the given node.
      */
     public int[] partitions(UUID nodeId) {
+        if (F.isEmpty(assignments))
+            return null;
+
         GridIntList parts = new GridIntList(assignments.size());
 
         for (int i = 0; i < assignments.size(); i++) {

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

@@ -132,19 +132,31 @@ public List<UUID> nodeIds() {
     }
 
     /** */
-    public FragmentMapping finalize(Supplier<List<UUID>> nodesSource) {
+    public FragmentMapping finalizeMapping(Supplier<List<UUID>> nodesSource) {
         if (colocationGroups.isEmpty())
             return this;
 
         List<ColocationGroup> colocationGroups = this.colocationGroups;
 
-        colocationGroups = Commons.transform(colocationGroups, ColocationGroup::finalaze);
+        colocationGroups = Commons.transform(colocationGroups, ColocationGroup::finalizeMapping);
         List<UUID> nodes = nodeIds(), nodes0 = nodes.isEmpty() ? nodesSource.get() : nodes;
         colocationGroups = Commons.transform(colocationGroups, g -> g.mapToNodes(nodes0));
 
         return new FragmentMapping(colocationGroups);
     }
 
+    /** */
+    public FragmentMapping filterByPartitions(int[] parts) throws ColocationMappingException {
+        List<ColocationGroup> colocationGroups = this.colocationGroups;
+
+        if (!F.isEmpty(parts) && colocationGroups.size() > 1)
+            throw new ColocationMappingException("Execution of non-collocated query with partition parameter is not possible");
+
+        colocationGroups = Commons.transform(colocationGroups, g -> g.filterByPartitions(parts));
+
+        return new FragmentMapping(colocationGroups);
+    }
+
     /** */
     public @NotNull ColocationGroup findGroup(long sourceId) {
         List<ColocationGroup> groups = colocationGroups.stream()

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

@@ -18,6 +18,7 @@
 package org.apache.ignite.internal.processors.query.calcite.prepare;
 
 import java.lang.reflect.Method;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Properties;
 import com.google.common.collect.Multimap;
@@ -159,14 +160,18 @@ public final class BaseQueryContext extends AbstractQueryContext {
     /** */
     private final boolean isLocal;
 
+    /** */
+    private final int[] parts;
+
     /**
      * Private constructor, used by a builder.
      */
     private BaseQueryContext(
         FrameworkConfig cfg,
         Context parentCtx,
         IgniteLogger log,
-        boolean isLocal
+        boolean isLocal,
+        int[] parts
     ) {
         super(Contexts.chain(parentCtx, cfg.getContext()));
 
@@ -177,6 +182,8 @@ private BaseQueryContext(
 
         this.isLocal = isLocal;
 
+        this.parts = parts;
+
         qryCancel = unwrap(GridQueryCancel.class);
 
         typeFactory = TYPE_FACTORY;
@@ -276,6 +283,14 @@ public boolean isLocal() {
         return isLocal;
     }
 
+    /** */
+    public int[] partitions() {
+        if (parts != null)
+            return Arrays.copyOf(parts, parts.length);
+
+        return null;
+    }
+
     /**
      * Query context builder.
      */
@@ -299,6 +314,9 @@ public static class Builder {
         /** */
         private boolean isLocal = false;
 
+        /** */
+        private int[] parts = null;
+
         /**
          * @param frameworkCfg Framework config.
          * @return Builder for chaining.
@@ -335,13 +353,24 @@ public Builder local(boolean isLocal) {
             return this;
         }
 
+        /**
+         * @param parts Array of partitions' numbers.
+         * @return Builder for chaining.
+         */
+        public Builder partitions(int[] parts) {
+            if (parts != null)
+                this.parts = Arrays.copyOf(parts, parts.length);
+
+            return this;
+        }
+
         /**
          * Builds planner context.
          *
          * @return Planner context.
          */
         public BaseQueryContext build() {
-            return new BaseQueryContext(frameworkCfg, parentCtx, log, isLocal);
+            return new BaseQueryContext(frameworkCfg, parentCtx, log, isLocal, parts);
         }
     }
 }

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

@@ -125,6 +125,11 @@ public Fragment attach(RelOptCluster cluster) {
         return root.getCluster() == cluster ? this : new Cloner(cluster).go(this);
     }
 
+    /** */
+    public Fragment filterByPartitions(int[] parts) throws ColocationMappingException {
+        return new Fragment(id, root, remotes, rootSer, mapping.filterByPartitions(parts));
+    }
+
     /**
      * Mapps the fragment to its data location.
      * @param ctx Planner context.
@@ -157,7 +162,7 @@ private FragmentMapping mapping(MappingQueryContext ctx, RelMetadataQuery mq, Su
                     .get(ThreadLocalRandom.current().nextInt(mapping.nodeIds().size()))).colocate(mapping);
             }
 
-            return mapping.finalize(nodesSource);
+            return mapping.finalizeMapping(nodesSource);
         }
         catch (NodeMappingException e) {
             throw new FragmentMappingException("Failed to calculate physical distribution", this, e.node(), e);