DRILL-4857: Maintain pruning status and populate ParquetGroupScan's entries field with only the selection root if no partition pruning was done.

close apache/drill#575

3 files changed, 53 insertions, 7 deletions

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/partition/PruneScanRule.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/partition/PruneScanRule.java
index 011c7834b..9a3ef96e5 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/partition/PruneScanRule.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/partition/PruneScanRule.java
@@ -61,6 +61,7 @@ import org.apache.drill.exec.record.VectorContainer;
 import org.apache.drill.exec.store.StoragePluginOptimizerRule;
 import org.apache.drill.exec.store.dfs.FormatSelection;
 import org.apache.drill.exec.store.dfs.MetadataContext;
+import org.apache.drill.exec.store.dfs.MetadataContext.PruneStatus;
 import org.apache.drill.exec.vector.NullableBitVector;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.plan.RelOptRule;
@@ -153,6 +154,12 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
     PartitionDescriptor descriptor = getPartitionDescriptor(settings, scanRel);
     final BufferAllocator allocator = optimizerContext.getAllocator();
 
+    final Object selection = getDrillTable(scanRel).getSelection();
+    MetadataContext metaContext = null;
+    if (selection instanceof FormatSelection) {
+         metaContext = ((FormatSelection)selection).getSelection().getMetaContext();
+    }
+
     RexNode condition = null;
     if (projectRel == null) {
       condition = filterRel.getCondition();
@@ -186,6 +193,7 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
     if (partitionColumnBitSet.isEmpty()) {
       logger.info("No partition columns are projected from the scan..continue. " +
           "Total pruning elapsed time: {} ms", totalPruningTime.elapsed(TimeUnit.MILLISECONDS));
+      setPruneStatus(metaContext, PruneStatus.NOT_PRUNED);
       return;
     }
 
@@ -207,6 +215,7 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
     if (pruneCondition == null) {
       logger.info("No conditions were found eligible for partition pruning." +
           "Total pruning elapsed time: {} ms", totalPruningTime.elapsed(TimeUnit.MILLISECONDS));
+      setPruneStatus(metaContext, PruneStatus.NOT_PRUNED);
       return;
     }
 
@@ -260,6 +269,7 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
             // materializePruneExpr logs it already
             logger.info("Total pruning elapsed time: {} ms",
                 totalPruningTime.elapsed(TimeUnit.MILLISECONDS));
+            setPruneStatus(metaContext, PruneStatus.NOT_PRUNED);
             return;
           }
         }
@@ -329,6 +339,8 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
       } catch (Exception e) {
         logger.warn("Exception while trying to prune partition.", e);
         logger.info("Total pruning elapsed time: {} ms", totalPruningTime.elapsed(TimeUnit.MILLISECONDS));
+
+        setPruneStatus(metaContext, PruneStatus.NOT_PRUNED);
         return; // continue without partition pruning
       } finally {
         container.clear();
@@ -407,11 +419,6 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
 
       }
 
-      final Object selection = getDrillTable(scanRel).getSelection();
-      MetadataContext metaContext = null;
-      if (selection instanceof FormatSelection) {
-           metaContext = ((FormatSelection)selection).getSelection().getMetaContext();
-      }
       RelNode inputRel = descriptor.supportsMetadataCachePruning() ?
           descriptor.createTableScan(newPartitions, cacheFileRoot, wasAllPartitionsPruned, metaContext) :
             descriptor.createTableScan(newPartitions, wasAllPartitionsPruned);
@@ -427,6 +434,8 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
         call.transformTo(newFilter);
       }
 
+      setPruneStatus(metaContext, PruneStatus.PRUNED);
+
     } catch (Exception e) {
       logger.warn("Exception while using the pruned partitions.", e);
     } finally {
@@ -520,4 +529,10 @@ public abstract class PruneScanRule extends StoragePluginOptimizerRule {
     return false;
   }
 
+  private static void setPruneStatus(MetadataContext metaContext, PruneStatus pruneStatus) {
+    if (metaContext != null) {
+      metaContext.setPruneStatus(pruneStatus);
+    }
+  }
+
 }
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/MetadataContext.java b/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/MetadataContext.java
index 17852ab5e..d115bee98 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/MetadataContext.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/MetadataContext.java
@@ -33,6 +33,14 @@ public class MetadataContext {
    */
   private Map<String, Boolean> dirModifCheckMap = Maps.newHashMap();
 
+  public enum PruneStatus {
+    NOT_STARTED,         // initial state
+    PRUNED,              // partitions were pruned
+    NOT_PRUNED           // partitions did not get pruned
+  }
+
+  private PruneStatus pruneStatus = PruneStatus.NOT_STARTED;
+
   public MetadataContext() {
   }
 
@@ -55,6 +63,18 @@ public class MetadataContext {
     dirModifCheckMap.clear();
   }
 
+  public void setPruneStatus(PruneStatus status) {
+    pruneStatus = status;
+  }
+
+  public boolean wasPruningStarted() {
+    return pruneStatus != PruneStatus.NOT_STARTED;
+  }
+
+  public boolean wasPruned() {
+    return pruneStatus == PruneStatus.PRUNED;
+  }
+
 }
 
 
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetGroupScan.java b/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetGroupScan.java
index 21227e410..8b255cfae 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetGroupScan.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetGroupScan.java
@@ -173,8 +173,18 @@ public class ParquetGroupScan extends AbstractFileGroupScan {
     final FileSelection fileSelection = expandIfNecessary(selection);
 
     this.entries = Lists.newArrayList();
-    for (String fileName : fileSelection.getFiles()) {
-      entries.add(new ReadEntryWithPath(fileName));
+    if (fileSelection.getMetaContext() != null &&
+        (fileSelection.getMetaContext().wasPruningStarted() &&
+        ! fileSelection.getMetaContext().wasPruned())) {
+      // if pruning was attempted and nothing was pruned, initialize the entries with just
+      // the selection root instead of the fully expanded list to reduce overhead. The fully
+      // expanded list is already stored as part of the fileSet.
+      // TODO: at some point we should examine whether the list of entries is absolutely needed.
+      entries.add(new ReadEntryWithPath(fileSelection.getSelectionRoot()));
+    } else {
+      for (String fileName : fileSelection.getFiles()) {
+        entries.add(new ReadEntryWithPath(fileName));
+      }
     }
 
     init(fileSelection.getMetaContext());
@@ -638,6 +648,7 @@ public class ParquetGroupScan extends AbstractFileGroupScan {
         cacheFileRoot, selection.wasAllPartitionsPruned());
 
     newSelection.setExpandedFully();
+    newSelection.setMetaContext(selection.getMetaContext());
     return newSelection;
   }