diff --git ql/src/java/org/apache/hadoop/hive/ql/optimizer/ColumnPrunerProcFactory.java ql/src/java/org/apache/hadoop/hive/ql/optimizer/ColumnPrunerProcFactory.java index 40c8939..2f98e66 100644 --- ql/src/java/org/apache/hadoop/hive/ql/optimizer/ColumnPrunerProcFactory.java +++ ql/src/java/org/apache/hadoop/hive/ql/optimizer/ColumnPrunerProcFactory.java @@ -61,6 +61,7 @@ import org.apache.hadoop.hive.ql.plan.AggregationDesc; import org.apache.hadoop.hive.ql.plan.ExprNodeColumnDesc; import org.apache.hadoop.hive.ql.plan.ExprNodeDesc; +import org.apache.hadoop.hive.ql.plan.ExprNodeDescUtils; import org.apache.hadoop.hive.ql.plan.GroupByDesc; import org.apache.hadoop.hive.ql.plan.JoinDesc; import org.apache.hadoop.hive.ql.plan.MapJoinDesc; @@ -695,12 +696,14 @@ private static void pruneReduceSinkOperator(boolean[] retainFlags, ReduceSinkOperator reduce, ColumnPrunerProcCtx cppCtx) throws SemanticException { ReduceSinkDesc reduceConf = reduce.getConf(); Map oldMap = reduce.getColumnExprMap(); + LOG.info("RS " + reduce.getIdentifier() + " oldColExprMap: " + oldMap); RowResolver oldRR = cppCtx.getOpToParseCtxMap().get(reduce).getRowResolver(); ArrayList signature = oldRR.getRowSchema().getSignature(); List valueColNames = reduceConf.getOutputValueColumnNames(); ArrayList newValueColNames = new ArrayList(); + List keyExprs = reduceConf.getKeyCols(); List valueExprs = reduceConf.getValueCols(); ArrayList newValueExprs = new ArrayList(); @@ -715,7 +718,12 @@ private static void pruneReduceSinkOperator(boolean[] retainFlags, } ColumnInfo colInfo = oldRR.getFieldMap(nm[0]).remove(nm[1]); oldRR.getInvRslvMap().remove(colInfo.getInternalName()); - oldMap.remove(outputCol); + + // Only remove the entry in colExprMap if the column is not a key + if (ExprNodeDescUtils.indexOf(outputColExpr, keyExprs) == -1) { + oldMap.remove(outputCol); + } + signature.remove(colInfo); } else { newValueColNames.add(outputCol); @@ -729,6 +737,7 @@ private static void pruneReduceSinkOperator(boolean[] retainFlags, .getFieldSchemasFromColumnList(reduceConf.getValueCols(), newValueColNames, 0, "")); reduceConf.setValueSerializeInfo(newValueTable); + LOG.info("RS " + reduce.getIdentifier() + " newColExprMap: " + oldMap); } /**