[drill] branch 1.20 updated: DRILL-8234: Register rules only from plugins used in the query (#2560)

[dz...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dzamo pushed a commit to branch 1.20
in repository https://gitbox.apache.org/repos/asf/drill.git


The following commit(s) were added to refs/heads/1.20 by this push:
     new 424d47de1f DRILL-8234: Register rules only from plugins used in the query (#2560)
424d47de1f is described below

commit 424d47de1fa848388c55589a8f4f939482d6954c
Author: Volodymyr Vysotskyi <vv...@gmail.com>
AuthorDate: Wed May 25 09:38:01 2022 +0300

    DRILL-8234: Register rules only from plugins used in the query (#2560)
    
    * DRILL-8234: Register rules only from plugins used in the query
    
    * DRILL-8234: Obtain plugins from query context
---
 .../planner/FileSystemPartitionDescriptor.java     |  3 +-
 .../planner/sql/handlers/DefaultSqlHandler.java    |  2 +-
 .../planner/sql/handlers/SqlHandlerConfig.java     | 52 ++++++++++++++++++----
 3 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/FileSystemPartitionDescriptor.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/FileSystemPartitionDescriptor.java
index 865ea04d49..5fb84708a7 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/FileSystemPartitionDescriptor.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/FileSystemPartitionDescriptor.java
@@ -31,7 +31,6 @@ import org.apache.commons.lang3.tuple.Pair;
 import org.apache.drill.common.util.GuavaUtils;
 import org.apache.drill.shaded.guava.com.google.common.base.Charsets;
 import org.apache.drill.shaded.guava.com.google.common.base.Preconditions;
-import org.apache.drill.shaded.guava.com.google.common.collect.ImmutableList;
 import org.apache.drill.shaded.guava.com.google.common.collect.Lists;
 
 import org.apache.calcite.adapter.enumerable.EnumerableTableScan;
@@ -251,7 +250,7 @@ public class FileSystemPartitionDescriptor extends AbstractPartitionDescriptor {
       DrillTranslatableTable newTable = new DrillTranslatableTable(dynamicDrillTable);
 
       RelOptTableImpl newOptTableImpl = RelOptTableImpl.create(relOptTable.getRelOptSchema(), relOptTable.getRowType(),
-          newTable, GuavaUtils.convertToUnshadedImmutableList(ImmutableList.of()));
+          newTable, GuavaUtils.convertToUnshadedImmutableList(relOptTable.getQualifiedName()));
 
       // return an EnumerableTableScan with fileSelection being part of digest of TableScan node.
       return DirPrunedEnumerableTableScan.create(scanRel.getCluster(), newOptTableImpl, newFileSelection.toString());
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java
index baa77b874b..48e8558e18 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java
@@ -364,7 +364,7 @@ public class DefaultSqlHandler extends AbstractSqlHandler {
   protected RelNode transform(PlannerType plannerType, PlannerPhase phase, RelNode input, RelTraitSet targetTraits,
       boolean log) {
     final Stopwatch watch = Stopwatch.createStarted();
-    final RuleSet rules = config.getRules(phase);
+    final RuleSet rules = config.getRules(phase, input);
     final RelTraitSet toTraits = targetTraits.simplify();
 
     final RelNode output;
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/SqlHandlerConfig.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/SqlHandlerConfig.java
index 9b7305e808..c2dd077fbe 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/SqlHandlerConfig.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/SqlHandlerConfig.java
@@ -17,16 +17,24 @@
  */
 package org.apache.drill.exec.planner.sql.handlers;
 
+import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Map.Entry;
+import java.util.List;
+import java.util.Optional;
 
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelShuttleImpl;
+import org.apache.calcite.rel.core.TableScan;
 import org.apache.calcite.tools.RuleSet;
+import org.apache.drill.common.util.function.CheckedSupplier;
 import org.apache.drill.exec.ops.QueryContext;
 import org.apache.drill.exec.planner.PlannerPhase;
+import org.apache.drill.exec.planner.common.DrillRelOptUtil;
+import org.apache.drill.exec.planner.logical.DrillTable;
+import org.apache.drill.exec.planner.sql.SchemaUtilites;
 import org.apache.drill.exec.planner.sql.conversion.SqlConverter;
 import org.apache.drill.exec.store.StoragePlugin;
-
-import org.apache.drill.shaded.guava.com.google.common.collect.Lists;
+import org.apache.drill.exec.store.StoragePluginRegistry;
 
 public class SqlHandlerConfig {
 
@@ -34,7 +42,6 @@ public class SqlHandlerConfig {
   private final SqlConverter converter;
 
   public SqlHandlerConfig(QueryContext context, SqlConverter converter) {
-    super();
     this.context = context;
     this.converter = converter;
   }
@@ -43,15 +50,42 @@ public class SqlHandlerConfig {
     return context;
   }
 
-  public RuleSet getRules(PlannerPhase phase) {
-    Collection<StoragePlugin> plugins = Lists.newArrayList();
-    for (Entry<String, StoragePlugin> k : context.getStorage()) {
-      plugins.add(k.getValue());
-    }
+  public RuleSet getRules(PlannerPhase phase, RelNode input) {
+    PluginsCollector pluginsCollector = new PluginsCollector(context.getStorage());
+    input.accept(pluginsCollector);
+
+    Collection<StoragePlugin> plugins = pluginsCollector.getPlugins();
     return phase.getRules(context, plugins);
   }
 
   public SqlConverter getConverter() {
     return converter;
   }
+
+  public static class PluginsCollector extends RelShuttleImpl {
+    private final List<StoragePlugin> plugins = new ArrayList<>();
+    private final StoragePluginRegistry storagePlugins;
+
+    public PluginsCollector(StoragePluginRegistry storagePlugins) {
+      this.storagePlugins = storagePlugins;
+    }
+
+    @Override
+    public RelNode visit(TableScan scan) {
+      String pluginName = SchemaUtilites.getSchemaPathAsList(
+        scan.getTable().getQualifiedName().iterator().next()).iterator().next();
+      CheckedSupplier<StoragePlugin, StoragePluginRegistry.PluginException> pluginsProvider =
+        () -> storagePlugins.getPlugin(pluginName);
+
+      StoragePlugin storagePlugin = Optional.ofNullable(DrillRelOptUtil.getDrillTable(scan))
+        .map(DrillTable::getPlugin)
+        .orElseGet(pluginsProvider);
+      plugins.add(storagePlugin);
+      return scan;
+    }
+
+    public List<StoragePlugin> getPlugins() {
+      return plugins;
+    }
+  }
 }