HIVE-29805: Write header/footer lines in HiveIgnoreKeyTextOutputForma…

ql/src/java/org/apache/hadoop/hive/ql/io/HiveIgnoreKeyTextOutputFormat.java

@@ -18,8 +18,10 @@
 
 package org.apache.hadoop.hive.ql.io;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.Properties;
 
 import org.apache.hadoop.fs.FileSystem;
@@ -76,9 +78,17 @@
     }
 
     final int finalRowSeparator = rowSeparator;
+    final int headerCount = getCount(tableProperties, serdeConstants.HEADER_COUNT);
+    final int footerCount = getCount(tableProperties, serdeConstants.FOOTER_COUNT);
+
     FileSystem fs = outPath.getFileSystem(jc);
     final OutputStream outStream = Utilities.createCompressedStream(jc,
     fs.create(outPath, progress), isCompressed);
+
+    if (headerCount > 0) {
+      outStream.write(buildHeader(tableProperties, rowSeparator, headerCount));
+    }
+
     return new RecordWriter() {
       @Override
       public void write(Writable r) throws IOException {
@@ -96,6 +106,11 @@
 
       @Override
       public void close(boolean abort) throws IOException {
+        if (!abort && footerCount > 0) {
+          for (int i = 0; i < footerCount; i++) {
+            outStream.write(finalRowSeparator);
+          }
+        }
         outStream.close();
       }
     };
@@ -130,4 +145,106 @@
         progress));
   }
 
+  private static int getCount(Properties tableProperties, String propertyName) {
+    String value = tableProperties.getProperty(propertyName, "0");
+    try {
+      return Integer.parseInt(value);
+    } catch (NumberFormatException e) {
+      return 0;
+    }
+  }
+
+  private static char getChar(Properties tableProperties, String propertyName, char defaultValue) {
+    String value = tableProperties.getProperty(propertyName);
+    if (value != null && !value.isEmpty()) {
+      return value.charAt(0);
+    }
+    return defaultValue;
+  }
+
+  private static String escapeQuotes(String field, char quote, char escape) {
+    if (quote == '\0' || field.indexOf(quote) < 0) {
+      return field;
+    }
+    StringBuilder sb = new StringBuilder(field.length() * 2);
+    for (int i = 0; i < field.length(); i++) {
+      char c = field.charAt(i);
+      if (c == quote) {
+        if (escape == quote) {
+          sb.append(quote);
+          sb.append(quote);
+        } else {
+          sb.append(escape);
+          sb.append(quote);
+        }
+      } else {
+        sb.append(c);
+      }
+    }
+    return sb.toString();
+  }
+
+  private static byte[] buildHeader(Properties tableProperties, int rowSeparator, int headerCount)
+      throws IOException {
+    char separator = getChar(tableProperties, "separatorChar",
+        getChar(tableProperties, serdeConstants.FIELD_DELIM, ','));
+    char quote = getChar(tableProperties, "quoteChar",
+        getChar(tableProperties, serdeConstants.QUOTE_CHAR, '"'));
+    char escape = getChar(tableProperties, "escapeChar",
+        getChar(tableProperties, serdeConstants.ESCAPE_CHAR, '"'));
+
+    // OpenCSVSerde quotes all data fields by default, so mirror that for the
+    // header unless the table explicitly disables it. For other text SerDes
+    // (e.g. LazySimpleSerDe), only quote a field when the character set requires it.
+    String serializationLib = tableProperties.getProperty("serialization.lib", "");
+    boolean openCsv = serializationLib.contains("OpenCSV");
+    String applyQuotesToAllProp = tableProperties.getProperty("applyQuotesToAll");
+    boolean applyQuotesToAll = applyQuotesToAllProp != null
+        ? Boolean.parseBoolean(applyQuotesToAllProp)
+        : openCsv;
+
+    String columns = tableProperties.getProperty(serdeConstants.LIST_COLUMNS);
+    ByteArrayOutputStream header = new ByteArrayOutputStream();
+    if (columns == null || columns.isEmpty()) {
+      // No column names are available; write empty header lines so that the
+      // configured number of header lines are skipped during reads.
+      for (int i = 0; i < headerCount; i++) {
+        header.write((byte) rowSeparator);
+      }
+      return header.toByteArray();
+    }
+
+    String[] colNames = columns.split(",");
+    for (int line = 0; line < headerCount; line++) {
+      if (line > 0) {
+        header.write((byte) rowSeparator);
+      }
+      for (int i = 0; i < colNames.length; i++) {
+        if (i > 0) {
+          header.write(String.valueOf(separator).getBytes(StandardCharsets.UTF_8));
+        }
+        if (quote != '\0' && (applyQuotesToAll
+            || needsQuoting(colNames[i], separator, quote, (char) rowSeparator))) {
+          header.write(String.valueOf(quote).getBytes(StandardCharsets.UTF_8));
+          header.write(escapeQuotes(colNames[i], quote, escape).getBytes(StandardCharsets.UTF_8));
+          header.write(String.valueOf(quote).getBytes(StandardCharsets.UTF_8));
+        } else {
+          header.write(colNames[i].getBytes(StandardCharsets.UTF_8));
+        }
+      }
+    }
+    header.write((byte) rowSeparator);
+    return header.toByteArray();
+  }
+
+  private static boolean needsQuoting(String field, char separator, char quote, char rowSeparator) {
+    for (int i = 0; i < field.length(); i++) {
+      char c = field.charAt(i);
+      if (c == separator || c == quote || c == rowSeparator) {
+        return true;
+      }
+    }
+    return false;
+  }
+
 }

ql/src/test/org/apache/hadoop/hive/ql/io/TestHiveIgnoreKeyTextOutputFormat.java

@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hive.ql.io;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hive.ql.exec.FileSinkOperator;
+import org.apache.hadoop.hive.serde.serdeConstants;
+import org.apache.hadoop.io.LongWritable;
+import org.apache.hadoop.io.Text;
+import org.apache.hadoop.io.WritableComparable;
+import org.apache.hadoop.io.Writable;
+import org.apache.hadoop.mapred.FileInputFormat;
+import org.apache.hadoop.mapred.InputSplit;
+import org.apache.hadoop.mapred.JobConf;
+import org.apache.hadoop.mapred.RecordReader;
+import org.apache.hadoop.mapred.Reporter;
+import org.apache.hadoop.util.Progressable;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Properties;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * Unittest for HiveIgnoreKeyTextOutputFormat header/footer writing.
+ */
+public class TestHiveIgnoreKeyTextOutputFormat {
+
+  private Configuration conf;
+  private JobConf job;
+  private FileSystem fileSystem;
+  private Path testDir;
+
+  @Before
+  public void setUp() throws IOException {
+    conf = new Configuration();
+    job = new JobConf(conf);
+    fileSystem = FileSystem.getLocal(conf);
+    testDir = new Path(System.getProperty("test.tmp.dir", System.getProperty(
+        "user.dir", new File(".").getAbsolutePath()))
+        + "/TestHiveIgnoreKeyTextOutputFormat");
+    fileSystem.delete(testDir, true);
+    fileSystem.mkdirs(testDir);
+  }
+
+  @After
+  public void tearDown() throws IOException {
+    fileSystem.delete(testDir, true);
+  }
+
+  /**
+   * Test that HiveIgnoreKeyTextOutputFormat writes a header and footer so that
+   * SkippingTextInputFormat does not lose data rows.
+   */
+  @Test
+  public void testHeaderAndFooterAreWritten() throws Exception {
+    Path outFile = new Path(testDir, "test.csv");
+
+    Properties tableProperties = new Properties();
+    tableProperties.setProperty(serdeConstants.LIST_COLUMNS, "a,b");
+    tableProperties.setProperty(serdeConstants.FIELD_DELIM, ",");
+    tableProperties.setProperty(serdeConstants.LINE_DELIM, "\n");
+    tableProperties.setProperty(serdeConstants.HEADER_COUNT, "1");
+    tableProperties.setProperty(serdeConstants.FOOTER_COUNT, "2");
+
+    HiveIgnoreKeyTextOutputFormat<WritableComparable, Writable> outputFormat =
+        new HiveIgnoreKeyTextOutputFormat<>();
+    FileSinkOperator.RecordWriter writer = outputFormat.getHiveRecordWriter(
+        job, outFile, Text.class, false, tableProperties, new Progressable() {
+          @Override
+          public void progress() {
+          }
+        });
+
+    writer.write(new Text("x,y"));
+    writer.write(new Text("a,b"));
+    writer.write(new Text("c,d"));
+    writer.close(false);
+
+    // Read the file back with header and footer skipping.
+    SkippingTextInputFormat inputFormat = new SkippingTextInputFormat();
+    inputFormat.configure(job, 1, 2);
+    FileInputFormat.setInputPaths(job, outFile);
+    InputSplit[] splits = inputFormat.getSplits(job, 2);
+
+    List<String> received = new ArrayList<>();
+    for (InputSplit split : splits) {
+      RecordReader<LongWritable, Text> reader =
+          inputFormat.getRecordReader(split, job, Reporter.NULL);
+      LongWritable key = reader.createKey();
+      Text value = reader.createValue();
+      while (reader.next(key, value)) {
+        received.add(value.toString());
+      }
+      reader.close();
+    }
+
+    assertEquals(Arrays.asList("x,y", "a,b", "c,d"), received);
+  }
+}

ql/src/test/results/clientpositive/llap/skip_header_footer_proj.q.out

@@ -30,6 +30,7 @@ POSTHOOK: query: SELECT * FROM hf1
 POSTHOOK: type: QUERY
 POSTHOOK: Input: default@hf1
 POSTHOOK: Output: hdfs://### HDFS PATH ###
+x	y
 a	b
 c	d
 PREHOOK: query: DROP TABLE IF EXISTS hf2
@@ -65,6 +66,8 @@ POSTHOOK: type: QUERY
 POSTHOOK: Input: default@hf2
 POSTHOOK: Output: hdfs://### HDFS PATH ###
 x	y
+a	b
+c	d
 PREHOOK: query: DROP TABLE IF EXISTS hf3
 PREHOOK: type: DROPTABLE
 PREHOOK: Output: database:default
@@ -97,6 +100,9 @@ POSTHOOK: query: SELECT * FROM hf3
 POSTHOOK: type: QUERY
 POSTHOOK: Input: default@hf3
 POSTHOOK: Output: hdfs://### HDFS PATH ###
+x	y
+a	b
+c	d
 PREHOOK: query: DROP TABLE IF EXISTS hf4
 PREHOOK: type: DROPTABLE
 PREHOOK: Output: database:default
@@ -129,3 +135,6 @@ POSTHOOK: query: SELECT * FROM hf4
 POSTHOOK: type: QUERY
 POSTHOOK: Input: default@hf4
 POSTHOOK: Output: hdfs://### HDFS PATH ###
+x	y
+a	b
+c	d