DRILL-1232: Fix unsupported exception while using 'where' clause with partition names

2 files changed, 12 insertions, 4 deletions

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/easy/EasyFormatPlugin.java b/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/easy/EasyFormatPlugin.java
index bdab07f17..bf8e301c2 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/easy/EasyFormatPlugin.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/store/dfs/easy/EasyFormatPlugin.java
@@ -18,6 +18,7 @@
 package org.apache.drill.exec.store.dfs.easy;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
 import java.util.regex.Matcher;
@@ -126,15 +127,20 @@ public abstract class EasyFormatPlugin<T extends FormatPluginConfig> implements
     if (columns == null || columns.size() == 0) {
       selectAllColumns = true;
     } else {
+      List<SchemaPath> newColumns = Lists.newArrayList();
       Pattern pattern = Pattern.compile(String.format("%s[0-9]+", partitionDesignator));
       for (SchemaPath column : columns) {
         Matcher m = pattern.matcher(column.getAsUnescapedPath());
         if (m.matches()) {
-          scan.getColumns().remove(column);
           selectedPartitionColumns.add(Integer.parseInt(column.getAsUnescapedPath().toString().substring(partitionDesignator.length())));
+        } else {
+          newColumns.add(column);
         }
       }
+      // Create a new sub scan object with the new set of columns;
+      scan = new EasySubScan(scan.getWorkUnits(), scan.getFormatPlugin(), newColumns, scan.getSelectionRoot());
     }
+
     int numParts = 0;
     for(FileWork work : scan.getWorkUnits()){
       readers.add(getRecordReader(context, work, scan.getColumns()));
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetScanBatchCreator.java b/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetScanBatchCreator.java
index ae73af365..5edd3c51d 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetScanBatchCreator.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/store/parquet/ParquetScanBatchCreator.java
@@ -61,7 +61,6 @@ public class ParquetScanBatchCreator implements BatchCreator<ParquetRowGroupScan
     Preconditions.checkArgument(children.isEmpty());
     String partitionDesignator = context.getConfig().getString(ExecConstants.FILESYSTEM_PARTITION_COLUMN_LABEL);
     List<SchemaPath> columns = rowGroupScan.getColumns();
-
     List<RecordReader> readers = Lists.newArrayList();
 
     List<String[]> partitionColumns = Lists.newArrayList();
@@ -71,17 +70,20 @@ public class ParquetScanBatchCreator implements BatchCreator<ParquetRowGroupScan
     if (columns == null || columns.size() == 0) {
       selectAllColumns = true;
     } else {
+      List<SchemaPath> newColums = Lists.newArrayList();
       Pattern pattern = Pattern.compile(String.format("%s[0-9]+", partitionDesignator));
       for (SchemaPath column : columns) {
         Matcher m = pattern.matcher(column.getAsUnescapedPath());
         if (m.matches()) {
-          columns.remove(column);
           selectedPartitionColumns.add(Integer.parseInt(column.getAsUnescapedPath().toString().substring(partitionDesignator.length())));
+        } else {
+          newColums.add(column);
         }
       }
+      // Create the new row group scan with the new columns
+      rowGroupScan = new ParquetRowGroupScan(rowGroupScan.getStorageEngine(), rowGroupScan.getRowGroupReadEntries(), newColums, rowGroupScan.getSelectionRoot());
     }
 
-
     FileSystem fs = rowGroupScan.getStorageEngine().getFileSystem().getUnderlying();
     Configuration conf = fs.getConf();
     conf.setBoolean(ENABLE_BYTES_READ_COUNTER, false);