content: /*
   * 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.carbondata.processing.loading.csvinput;
  
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.FileOutputStream;
  import java.io.IOException;
  
  import junit.framework.TestCase;
  import org.junit.Assert;
  import org.junit.Test;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.fs.Path;
  import org.apache.hadoop.io.NullWritable;
  import org.apache.hadoop.io.compress.BZip2Codec;
  import org.apache.hadoop.io.compress.CompressionOutputStream;
  import org.apache.hadoop.io.compress.GzipCodec;
  import org.apache.hadoop.io.compress.Lz4Codec;
  import org.apache.hadoop.io.compress.SnappyCodec;
  import org.apache.hadoop.mapreduce.Job;
  import org.apache.hadoop.mapreduce.Mapper;
  import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
  import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
  
  public class CSVInputFormatTest extends TestCase {
  
    /**
     * generate compressed files, no need to call this method.
     * @throws Exception
     */
    public void generateCompressFiles() throws Exception {
      String pwd = new File("src/test/resources/csv").getCanonicalPath();
      String inputFile = pwd + "/data.csv";
      FileInputStream input = new FileInputStream(inputFile);
      Configuration conf = new Configuration();
  
      // .gz
      String outputFile = pwd + "/data.csv.gz";
      FileOutputStream output = new FileOutputStream(outputFile);
      GzipCodec gzip = new GzipCodec();
      gzip.setConf(conf);
      CompressionOutputStream outputStream = gzip.createOutputStream(output);
      int i = -1;
      while ((i = input.read()) != -1) {
        outputStream.write(i);
      }
      outputStream.close();
      input.close();
  
      // .bz2
      input = new FileInputStream(inputFile);
      outputFile = pwd + "/data.csv.bz2";
      output = new FileOutputStream(outputFile);
      BZip2Codec bzip2 = new BZip2Codec();
      bzip2.setConf(conf);
      outputStream = bzip2.createOutputStream(output);
      i = -1;
      while ((i = input.read()) != -1) {
        outputStream.write(i);
      }
      outputStream.close();
      input.close();
  
      // .snappy
      input = new FileInputStream(inputFile);
      outputFile = pwd + "/data.csv.snappy";
      output = new FileOutputStream(outputFile);
      SnappyCodec snappy = new SnappyCodec();
      snappy.setConf(conf);
      outputStream = snappy.createOutputStream(output);
      i = -1;
      while ((i = input.read()) != -1) {
        outputStream.write(i);
      }
      outputStream.close();
      input.close();
  
      //.lz4
      input = new FileInputStream(inputFile);
      outputFile = pwd + "/data.csv.lz4";
      output = new FileOutputStream(outputFile);
      Lz4Codec lz4 = new Lz4Codec();
      lz4.setConf(conf);
      outputStream = lz4.createOutputStream(output);
      i = -1;
      while ((i = input.read()) != -1) {
        outputStream.write(i);
      }
      outputStream.close();
      input.close();
  
    }
  
    /**
     * CSVCheckMapper check the content of csv files.
     */
    public static class CSVCheckMapper extends Mapper<NullWritable, StringArrayWritable, NullWritable,
        NullWritable> {
      @Override
      protected void map(NullWritable key, StringArrayWritable value, Context context)
          throws IOException, InterruptedException {
        String[] columns = value.get();
        int id = Integer.parseInt(columns[0]);
        int salary = Integer.parseInt(columns[6]);
        Assert.assertEquals(id - 1, salary - 15000);
      }
    }
  
    /**
     * test read csv files
     * @throws Exception
     */
    @Test public void testReadCSVFiles() throws Exception{
      Configuration conf = new Configuration();
      prepareConf(conf);
      conf.setBoolean(CSVInputFormat.HEADER_PRESENT, true);
      File output = new File("target/output_CSVInputFormatTest");
      conf.set("mapreduce.cluster.local.dir", output.getCanonicalPath());
      Job job = Job.getInstance(conf, "CSVInputFormat_normal");
      job.setJarByClass(CSVInputFormatTest.class);
      job.setMapperClass(CSVCheckMapper.class);
      job.setNumReduceTasks(0);
      job.setInputFormatClass(CSVInputFormat.class);
  
      String inputFolder = new File("src/test/resources/csv").getCanonicalPath();
      FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "data.csv"));
      FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "data.csv.bz2"));
      FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "data.csv.gz"));
      // FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "data.csv.lz4"));
      // FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "data.csv.snappy"));
  
      deleteOutput(output);
      FileOutputFormat.setOutputPath(job, new Path(output.getCanonicalPath()));
  
      Assert.assertTrue(job.waitForCompletion(true));
    }
  
    /**
     * test read csv files encoded as UTF-8 with BOM
     * @throws Exception
     */
    @Test public void testReadCSVFilesWithBOM() throws Exception{
  
      Configuration conf = new Configuration();
      prepareConf(conf);
      conf.setBoolean(CSVInputFormat.HEADER_PRESENT, false);
      File output = new File("target/output_CSVInputFormatTest_bom");
      conf.set("mapreduce.cluster.local.dir", output.getCanonicalPath());
      Job job = Job.getInstance(conf, "CSVInputFormat_normal_bom");
      job.setJarByClass(CSVInputFormatTest.class);
      job.setMapperClass(CSVCheckMapper.class);
      job.setNumReduceTasks(0);
      job.setInputFormatClass(CSVInputFormat.class);
  
      String inputFolder = new File("src/test/resources/csv").getCanonicalPath();
      FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "csv_with_bom.csv"));
      FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "csv_with_bom.csv.bz2"));
      FileInputFormat.addInputPath(job, new Path(inputFolder + File.separator + "csv_with_bom.csv.gz"));
  
      deleteOutput(output);
      FileOutputFormat.setOutputPath(job, new Path(output.getCanonicalPath()));
  
      Assert.assertTrue(job.waitForCompletion(true));
      deleteOutput(output);
    }
  
    private void prepareConf(Configuration conf) {
      conf.set(CSVInputFormat.MAX_COLUMNS, "10");
      conf.set(CSVInputFormat.NUMBER_OF_COLUMNS, "7");
    }
  
    private void deleteOutput(File output) {
      if (output.exists()) {
        if (output.isDirectory()) {
          for(File file : output.listFiles()) {
            deleteOutput(file);
          }
          output.delete();
        } else {
          output.delete();
        }
      }
    }
  }
  