HIVE-29262: [Addendum] Fix VectorPTFGroupBatches partition key copy when partition and order columns overlap

ql/src/java/org/apache/hadoop/hive/ql/exec/vector/ptf/VectorPTFGroupBatches.java

@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.stream.Collectors;
 
+import org.apache.commons.lang3.ArrayUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hive.common.type.HiveIntervalDayTime;
 import org.apache.hadoop.hive.ql.exec.PTFPartition;
@@ -64,6 +65,7 @@
   private int[] bufferedColumnMap;
   private int[] orderColumnMap;
   private int[] keyWithoutOrderColumnMap;
+  private int[] partitionKeyIndices;
 
   private int spillLimitBufferedBatchCount;
   private String spillLocalDirs;
@@ -287,15 +289,23 @@
     initValueCache(hconf);
   }
 
   public void init(VectorPTFEvaluatorBase[] evaluators, int[] outputProjectionColumnMap,
       int[] bufferedColumnMap, TypeInfo[] bufferedTypeInfos, int[] orderColumnMap,
-      int[] keyWithoutOrderColumnMap, VectorizedRowBatch overflowBatch) {
+      int[] keyWithoutOrderColumnMap, int[] partitionColumnMap, VectorizedRowBatch overflowBatch) {
     this.evaluators = evaluators;
     this.outputProjectionColumnMap = outputProjectionColumnMap;
     this.bufferedColumnMap = bufferedColumnMap;
     this.orderColumnMap = orderColumnMap;
     this.keyWithoutOrderColumnMap = keyWithoutOrderColumnMap;
 
+    partitionKeyIndices = new int[keyWithoutOrderColumnMap.length];
+    for (int i = 0; i < keyWithoutOrderColumnMap.length; i++) {
+      partitionKeyIndices[i] = ArrayUtils.indexOf(partitionColumnMap, keyWithoutOrderColumnMap[i]);
+      Preconditions.checkState(partitionKeyIndices[i] != ArrayUtils.INDEX_NOT_FOUND,
+          "Key input column %s not found among partition columns %s",
+          keyWithoutOrderColumnMap[i], Arrays.toString(partitionColumnMap));
+    }
+
     this.overflowBatch = overflowBatch;
     bufferedBatches = new ArrayList<BufferedVectorizedRowBatch>(0);
 
@@ -775,14 +785,19 @@
     }
   }
 
+  /**
+   * Sets the partition key values as repeating columns in the overflowBatch. Key columns are in
+   * output projection order, partitionKey values in PARTITION BY order; partitionKeyIndices maps
+   * between the two.
+   */
   private void copyPartitionColumnToOverflow(Object[] partitionKey) {
     // Set partition column in overflowBatch.
     // We can set by ref since our last batch is held by us.
     final int keyInputColumnCount = keyWithoutOrderColumnMap.length;
     for (int i = 0; i < keyInputColumnCount; i++) {
       final int keyColumnNum = keyWithoutOrderColumnMap[i];
       Preconditions.checkState(overflowBatch.cols[keyColumnNum] != null);
-      setRepeatingColumn(partitionKey[i], overflowBatch, keyColumnNum);
+      setRepeatingColumn(partitionKey[partitionKeyIndices[i]], overflowBatch, keyColumnNum);
     }
   }
 

ql/src/java/org/apache/hadoop/hive/ql/exec/vector/ptf/VectorPTFOperator.java

@@ -349,6 +349,7 @@ protected void initializeOp(Configuration hconf) throws HiveException {
         bufferedTypeInfos,
         orderColumnMap,
         keyWithoutOrderColumnMap,
+        vectorPTFInfo.getPartitionColumnMap(),
         overflowBatch);
 
     isFirstPartition = true;

ql/src/java/org/apache/hadoop/hive/ql/optimizer/physical/Vectorizer.java

@@ -40,7 +40,6 @@
 import java.util.TreeSet;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.apache.commons.lang3.StringUtils;
@@ -5077,59 +5076,6 @@ private static void createVectorPTFDesc(Operator<? extends OperatorDesc> ptfOp,
         vectorizedPTFMaxMemoryBufferingBatchCount);
   }
 
-  /**
-   * Reorders partitionColumnMap and partitionColumnVectorTypes in-place so that projected
-   * partition-only columns come first in SELECT output order. Any partition columns not found in
-   * the output are appended at the end in their original plan order.
-   */
-  private static void reorderPartitionColumnsToMatchOutputOrder(List<ColumnInfo> outputSignature,
-      int evaluatorCount, int[] outputColumnProjectionMap, int[] orderColumnMap,
-      ExprNodeDesc[] partitionExprNodeDescs, int[] partitionColumnMap,
-      Type[] partitionColumnVectorTypes) {
-    final int count = partitionColumnMap.length;
-    final int[] orderedMap = new int[count];
-    final Type[] orderedTypes = new Type[count];
-    final boolean[] placed = new boolean[count];
-
-    int idx = 0;
-    final int outputSize = outputSignature.size();
-
-    for (int outputIdx = evaluatorCount; outputIdx < outputSize && idx < count; outputIdx++) {
-      final int outputColumn = outputColumnProjectionMap[outputIdx];
-      final String colName = outputSignature.get(outputIdx).getInternalName();
-      int matchedPartitionIdx = IntStream.range(0, count)
-          .filter(p -> !placed[p])
-          .filter(p -> partitionExprNodeDescs[p] instanceof ExprNodeColumnDesc colDesc &&
-              colDesc.getColumn().equals(colName))
-          .findFirst()
-          .orElse(-1);
-
-      if (matchedPartitionIdx == -1) {
-        matchedPartitionIdx = IntStream.range(0, count)
-            .filter(p -> !placed[p] && partitionColumnMap[p] == outputColumn)
-            .findFirst()
-            .orElse(-1);
-      }
-
-      if (matchedPartitionIdx != -1 && !ArrayUtils.contains(orderColumnMap, outputColumn)) {
-        orderedMap[idx] = outputColumn;
-        orderedTypes[idx] = partitionColumnVectorTypes[matchedPartitionIdx];
-        placed[matchedPartitionIdx] = true;
-        idx++;
-      }
-    }
-
-    for (int p = 0; p < count; p++) {
-      if (!placed[p]) {
-        orderedMap[idx] = partitionColumnMap[p];
-        orderedTypes[idx] = partitionColumnVectorTypes[p];
-        idx++;
-      }
-    }
-    System.arraycopy(orderedMap, 0, partitionColumnMap, 0, count);
-    System.arraycopy(orderedTypes, 0, partitionColumnVectorTypes, 0, count);
-  }
-
   private static void determineKeyAndNonKeyInputColumnMap(int[] outputColumnProjectionMap,
       boolean isPartitionOrderBy, int[] orderColumnMap, int[] partitionColumnMap,
       int evaluatorCount, ArrayList<Integer> keyInputColumns,
@@ -5244,11 +5190,6 @@ private static VectorPTFInfo createVectorPTFInfo(Operator<? extends OperatorDesc
     int[] keyInputColumnMap = ArrayUtils.toPrimitive(keyInputColumns.toArray(new Integer[0]));
     int[] nonKeyInputColumnMap = ArrayUtils.toPrimitive(nonKeyInputColumns.toArray(new Integer[0]));
 
-    if (isPartitionOrderBy && partitionKeyCount > 1) {
-      reorderPartitionColumnsToMatchOutputOrder(outputSignature, evaluatorCount, outputColumnProjectionMap,
-          orderColumnMap, partitionExprNodeDescs, partitionColumnMap, partitionColumnVectorTypes);
-    }
-
     VectorExpression[][] evaluatorInputExpressions = new VectorExpression[evaluatorCount][];
     Type[][] evaluatorInputColumnVectorTypes = new Type[evaluatorCount][];
     for (int i = 0; i < evaluatorCount; i++) {

ql/src/test/org/apache/hadoop/hive/ql/exec/vector/ptf/TestVectorPTFGroupBatches.java

@@ -588,6 +588,7 @@
         /* bufferedTypeInfos */ new TypeInfo[] { getTypeInfo("int"), getTypeInfo("string") },
         /* orderColumnMap */ new int[] { 1 }, // p_date
         /* keyWithoutOrderColumnMap */ new int[] { 0 }, // p_mfgr
+        /* partitionColumnMap */ new int[] { 0 },
         getFakeOperator().setupOverflowBatch(3, new String[] { "bigint", "bigint" },
             outputProjectionColumnMap, outputTypeInfos));
   }

ql/src/test/queries/clientpositive/vector_ptf_spill_partition_order_overlap.q

@@ -0,0 +1,39 @@
+SET hive.vectorized.execution.enabled=true;
+SET hive.vectorized.execution.reduce.enabled=true;
+SET hive.vectorized.execution.ptf.enabled=true;
+SET hive.vectorized.testing.reducer.batch.size=2;
+SET hive.vectorized.ptf.max.memory.buffering.batch.count=1;
+
+SET hive.fetch.task.conversion=none;
+
+CREATE TABLE t1 (
+  dept STRING,
+  region BIGINT,
+  emp_id BIGINT,
+  salary DOUBLE
+) STORED AS ORC;
+
+INSERT INTO t1 values
+  ('engineering', 10, 1, 50000.0),
+  ('engineering', 10, 2, 55000.0),
+  ('engineering', 10, 3, 60000.0),
+  ('engineering', 10, 4, 45000.0),
+  ('engineering', 10, 5, 70000.0);
+
+EXPLAIN VECTORIZATION DETAIL
+SELECT dept, region, emp_id,
+  SUM(salary) OVER (
+    PARTITION BY dept, region
+    ORDER BY dept
+    ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+  ) AS total
+FROM t1;
+
+SELECT dept, region, emp_id,
+  SUM(salary) OVER (
+    PARTITION BY dept, region
+    ORDER BY dept
+    ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+  ) AS total
+FROM t1
+order by emp_id;