HIVE-29227: Apply masking in memory instead of rewriting the file in QOutProcessor#maskPatterns

common/src/java/org/apache/hadoop/hive/common/io/CachingPrintStream.java

@@ -20,6 +20,7 @@
 
 import java.io.FileNotFoundException;
 import java.io.OutputStream;
+import java.io.PrintStream;
 import java.io.UnsupportedEncodingException;
 import java.util.ArrayList;
 import java.util.List;
@@ -43,7 +44,12 @@ public CachingPrintStream(OutputStream out) {
   @Override
   public void println(String out) {
     output.add(out);
-    super.println(out);
+    // Delegate to a wrapped PrintStream so downstream transforms (e.g. QTestFetchConverter) apply.
+    if (this.out instanceof PrintStream ps && ps != this) {
+      ps.println(out);
+    } else {
+      super.println(out);
+    }
   }
 
   @Override

common/src/java/org/apache/hadoop/hive/common/io/FetchConverter.java

@@ -19,6 +19,7 @@
 package org.apache.hadoop.hive.common.io;
 
 import java.io.OutputStream;
+import java.io.PrintStream;
 import java.io.UnsupportedEncodingException;
 
 public abstract class FetchConverter extends SessionStream implements FetchCallback {
@@ -48,8 +49,14 @@ public void println(String out) {
     }
   }
 
-  protected final void printDirect(String out) {
-    super.println(out);
+  protected final void printDirect(String line) {
+    // Delegate to a wrapped PrintStream so downstream transforms (e.g. QTestFetchConverter) apply
+    // when this converter is stacked above them (sort/hash then mask).
+    if (this.out instanceof PrintStream ps && ps != this) {
+      ps.println(line);
+    } else {
+      super.println(line);
+    }
   }
 
   protected final boolean byPass() {

common/src/java/org/apache/hadoop/hive/common/io/QTestFetchConverter.java

@@ -43,7 +43,34 @@ public QTestFetchConverter(OutputStream out, boolean autoFlush, String encoding,
 
   @Override
   public void println(String str) {
-    inner.println(transformation.apply(str));
+    // PREHOOK/POSTHOOK may send multiple lines separated by \n in one println, mask each
+    // separately. Query fetch rows must stay intact so SORT_QUERY_RESULTS keeps row order.
+    if (str.indexOf('\n') >= 0 && shouldSplitMultiline()) {
+      int start = 0;
+      for (int i = 0; i < str.length(); i++) {
+        if (str.charAt(i) == '\n') {
+          writeTransformedLine(str.substring(start, i));
+          start = i + 1;
+        }
+      }
+      writeTransformedLine(str.substring(start));
+      return;
+    }
+    writeTransformedLine(str);
+  }
+
+  private boolean shouldSplitMultiline() {
+    if (inner instanceof FetchConverter fetchConverter) {
+      return !(fetchConverter.queryfound && fetchConverter.fetchStarted);
+    }
+    return true;
+  }
+
+  private void writeTransformedLine(String line) {
+    String transformed = transformation.apply(line);
+    if (transformed != null) {
+      inner.println(transformed);
+    }
   }
 
   @Override

itests/util/src/main/java/org/apache/hadoop/hive/cli/control/CoreBeeLineDriver.java

@@ -21,6 +21,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -35,6 +36,7 @@
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.hive.conf.HiveConf;
+import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
 import org.apache.hadoop.hive.conf.HiveConfUtil;
 import org.apache.hadoop.hive.ql.QTestProcessExecResult;
 import org.apache.hadoop.hive.ql.QTestUtil;
@@ -227,6 +229,11 @@ private void runTest(QFile qFile, List<Callable<Void>> preCommands) throws Excep
       long startTime = System.currentTimeMillis();
       System.err.println(">>> STARTED " + qFile.getName());
 
+      qOutProcessor.initMasks(FileUtils.readFileToString(qFile.getInputFile(), StandardCharsets.UTF_8));
+      setupAdditionalPartialMasks();
+      QOutProcessor.MaskingFoldState maskingFoldState = new QOutProcessor.MaskingFoldState();
+      beeLineClient.setOutputMasking(qOutProcessor, maskingFoldState);
+
       beeLineClient.execute(qFile, preCommands);
 
       long queryEndTime = System.currentTimeMillis();
@@ -239,7 +246,6 @@ private void runTest(QFile qFile, List<Callable<Void>> preCommands) throws Excep
           + "ms");
 
       if (!overwrite) {
-        qOutProcessor.maskPatterns(qFile.getOutputFile().getPath());
         QTestProcessExecResult result = qFile.compareResults();
 
         long compareEndTime = System.currentTimeMillis();
@@ -259,13 +265,40 @@ private void runTest(QFile qFile, List<Callable<Void>> preCommands) throws Excep
         qFile.overwriteResults();
         System.err.println(">>> PASSED " + qFile.getName());
       }
+      resetAdditionalPartialMasks();
     } catch (Exception e) {
+      resetAdditionalPartialMasks();
       throw new Exception("Exception running or analyzing the results of the query file: " + qFile
           + "\n" + qFile.getDebugHint(), e);
     }
 
   }
 
+  private void setupAdditionalPartialMasks() {
+    if (miniHS2 == null) {
+      return;
+    }
+    HiveConf conf = miniHS2.getHiveConf();
+    String patternStr = HiveConf.getVar(conf, ConfVars.HIVE_ADDITIONAL_PARTIAL_MASKS_PATTERN);
+    String replacementStr = HiveConf.getVar(conf, ConfVars.HIVE_ADDITIONAL_PARTIAL_MASKS_REPLACEMENT_TEXT);
+    if (patternStr != null && replacementStr != null && !replacementStr.isEmpty()
+        && !patternStr.isEmpty()) {
+      String[] patterns = patternStr.split(",");
+      String[] replacements = replacementStr.split(",");
+      if (patterns.length != replacements.length) {
+        throw new RuntimeException("Count mismatch for additional partial masks and their replacements");
+      }
+      for (int i = 0; i < patterns.length; i++) {
+        qOutProcessor.addPatternWithMaskComment(patterns[i],
+            String.format("### %s ###", replacements[i]));
+      }
+    }
+  }
+
+  private void resetAdditionalPartialMasks() {
+    qOutProcessor.resetPatternwithMaskComments();
+  }
+
   public void runTest(QFile qFile) throws Exception {
     runTest(qFile, null);
   }

itests/util/src/main/java/org/apache/hadoop/hive/cli/control/CoreCliDriver.java

@@ -109,6 +109,7 @@ public void runTest(String testName, String fname, String fpath) {
       System.err.println("Begin query: " + fname);
 
       qt.setInputFile(fpath);
+      setupAdditionalPartialMasks();
       qt.cliInit();
 
       try {
@@ -118,7 +119,6 @@ public void runTest(String testName, String fname, String fpath) {
         qt.failedQuery(e.getCause(), e.getResponseCode(), fname, QTestUtil.DEBUG_HINT);
       }
 
-      setupAdditionalPartialMasks();
       QTestProcessExecResult result = qt.checkCliDriverResults();
       resetAdditionalPartialMasks();
       if (result.getReturnCode() != 0) {

itests/util/src/main/java/org/apache/hadoop/hive/ql/QOutProcessor.java

@@ -18,20 +18,18 @@
 package org.apache.hadoop.hive.ql;
 
 import java.io.BufferedReader;
-import java.io.BufferedWriter;
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileOutputStream;
-import java.io.InputStreamReader;
-import java.io.OutputStreamWriter;
+import java.io.IOException;
+import java.io.StringReader;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
-import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.hadoop.hive.ql.QTestMiniClusters.FsType;
@@ -91,6 +89,11 @@ public String get() {
     }
   }
 
+  public static final class MaskingFoldState {
+    private boolean lastWasMasked;
+    private boolean lastWasVertexKilled;
+  }
+
   private final Pattern[] planMask = toPattern(new String[] {
       ".*[.][.][.] [0-9]* more.*",
       "pk_-?[0-9]*_[0-9]*_[0-9]*",
@@ -184,55 +187,73 @@ private Pattern[] toPattern(String[] patternStrs) {
     return patterns;
   }
 
-  public void maskPatterns(String fname) throws Exception {
-
-    String line;
-    BufferedReader in;
-    BufferedWriter out;
-
-    File file = new File(fname);
-    File fileOrig = new File(fname + ".orig");
-    FileUtils.copyFile(file, fileOrig);
-
-    in = new BufferedReader(new InputStreamReader(new FileInputStream(fileOrig), "UTF-8"));
-    out = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(file), "UTF-8"));
-
-    boolean lastWasMasked = false;
-    boolean lastWasVertexKilled = false;
+  /**
+   * Applies {@link #processLine(String)} and consecutive-line folding.
+   *
+   * @return the line to emit, or {@code null} if the line should be suppressed
+   */
+  public String maskAndFoldLine(String line, MaskingFoldState state) {
+    LineProcessingResult result = processLine(line);
+
+    if (result.line.equals(MASK_PATTERN)) {
+      // We're folding multiple masked lines into one.
+      if (state.lastWasMasked) {
+        state.lastWasVertexKilled = false;
+        return null;
+      }
+      state.lastWasMasked = true;
+      state.lastWasVertexKilled = false;
+      return result.line;
+    }
+    if (result.line.equals(MASKED_VERTEX_KILLED_PATTERN)) {
+      // Deduplicate consecutive standalone vertex-killed lines — the number of sibling
+      // vertices still alive when the kill propagates is non-deterministic.
+      if (state.lastWasVertexKilled) {
+        state.lastWasMasked = false;
+        return null;
+      }
+      state.lastWasVertexKilled = true;
+      state.lastWasMasked = false;
+      return result.line;
+    }
+    state.lastWasMasked = false;
+    state.lastWasVertexKilled = false;
+    return result.line;
+  }
 
-    while (null != (line = in.readLine())) {
-      LineProcessingResult result = processLine(line);
+  public List<String> maskLines(List<String> lines) {
+    MaskingFoldState state = new MaskingFoldState();
+    List<String> masked = new ArrayList<>(lines.size());
+    for (String line : lines) {
+      String folded = maskAndFoldLine(line, state);
+      if (folded != null) {
+        masked.add(folded);
+      }
+    }
+    return masked;
+  }
 
-      if (result.line.equals(MASK_PATTERN)) {
-        // We're folding multiple masked lines into one.
-        if (!lastWasMasked) {
-          out.write(result.line);
-          out.write("\n");
-          lastWasMasked = true;
-          result.partialMaskWasMatched = false;
-        }
-        lastWasVertexKilled = false;
-      } else if (result.line.equals(MASKED_VERTEX_KILLED_PATTERN)) {
-        // Deduplicate consecutive standalone vertex-killed lines — the number of sibling
-        // vertices still alive when the kill propagates is non-deterministic.
-        if (!lastWasVertexKilled) {
-          out.write(result.line);
-          out.write("\n");
-          lastWasVertexKilled = true;
-        }
-        lastWasMasked = false;
-        result.partialMaskWasMatched = false;
-      } else {
-        out.write(result.line);
-        out.write("\n");
-        lastWasMasked = false;
-        lastWasVertexKilled = false;
-        result.partialMaskWasMatched = false;
+  /** Applies masking and folding to multiline text */
+  public String maskContent(String content) throws IOException {
+    List<String> lines = new ArrayList<>();
+    try (BufferedReader reader = new BufferedReader(new StringReader(content))) {
+      String line;
+      while ((line = reader.readLine()) != null) {
+        lines.add(line);
       }
     }
+    StringBuilder out = new StringBuilder();
+    for (String line : maskLines(lines)) {
+      out.append(line).append('\n');
+    }
+    return out.toString();
+  }
 
-    in.close();
-    out.close();
+  /** Applies {@link #maskContent(String)} to a q test output file in place. */
+  public void maskOutputFile(String fname) throws IOException {
+    Path path = Path.of(fname);
+    String masked = maskContent(Files.readString(path, StandardCharsets.UTF_8));
+    Files.writeString(path, masked, StandardCharsets.UTF_8);
   }
 
   LineProcessingResult processLine(String line) {

itests/util/src/main/java/org/apache/hadoop/hive/ql/QTestResultProcessor.java

@@ -93,7 +93,7 @@ public void setOutputs(CliSessionState ss, OutputStream fo) throws Exception {
     // specified below. This ensures the behavior remains the same as before this
     // refactoring.
     // It would be better to throw an error than silently pick one and ignore the
-    // rest but it is out of the scope of the current change. 
+    // rest but it is out of the scope of the current change.
     if (operations.contains(Operation.SORT)) {
       ss.out = new SortPrintStream(fo, "UTF-8");
     } else if (operations.contains(Operation.HASH)) {

itests/util/src/main/java/org/apache/hadoop/hive/ql/QTestUtil.java

@@ -691,7 +691,7 @@ private void setSessionOutputs(CliSessionState ss, File outf) throws Exception {
     ss.out = new QTestFetchConverter(ss.out, false, "UTF-8", line -> {
       notifyOutputLine(line);
       if (qOutProcessor != null) {
-        // ensure that the masking is done before the sorting of the query results
+        // Mask before sort/hash (HIVE-29226); folding runs in maskOutputFile().
         return qOutProcessor.processLine(line).get();
       }
       return line;
@@ -1014,7 +1014,8 @@ public QTestProcessExecResult checkCliDriverResults() throws Exception {
     String outFileName = outPath(outDir, tname + outFileExtension);
 
     File f = new File(logDir, tname + outFileExtension);
-    qOutProcessor.maskPatterns(f.getPath());
+
+    qOutProcessor.maskOutputFile(f.getPath());
 
     if (QTestSystemProperties.shouldOverwriteResults()) {
       qTestResultProcessor.overwriteResults(f.getPath(), outFileName);