contrib/piggybank/java/src/test/java/org/apache/pig/piggybank/test/storage/TestMultiStorage.java

/*
 * 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.pig.piggybank.test.storage;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.util.List;
import java.util.Map;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.FileUtil;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.PathFilter;
import org.apache.pig.PigServer;
import org.apache.pig.backend.executionengine.ExecJob;
import org.apache.pig.test.MiniGenericCluster;
import org.apache.pig.test.Util;
import org.apache.pig.tools.pigstats.JobStats;
import org.apache.pig.tools.pigstats.OutputStats;
import org.apache.pig.tools.pigstats.PigStats;
import org.apache.pig.tools.pigstats.mapreduce.SimplePigStats;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.Test;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class TestMultiStorage {
  private static final String INPUT_FILE = "MultiStorageInput.txt";

  private PigServer pigServer;
  private PigServer pigServerLocal;

  private static MiniGenericCluster cluster = MiniGenericCluster.buildCluster();

  public TestMultiStorage() throws Exception {
    pigServer = new PigServer(cluster.getExecType(), cluster.getProperties());
    pigServerLocal = new PigServer(Util.getLocalTestMode());
  }

  public static final PathFilter hiddenPathFilter = new PathFilter() {
    public boolean accept(Path p) {
      String name = p.getName();
      return !name.startsWith("_") && !name.startsWith(".");
    }
  };

  private void createFile() throws IOException {
    PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
    w.println("100\tapple\taaa1");
    w.println("200\torange\tbbb1");
    w.println("300\tstrawberry\tccc1");

    w.println("101\tapple\taaa2");
    w.println("201\torange\tbbb2");
    w.println("301\tstrawberry\tccc2");

    w.println("102\tapple\taaa3");
    w.println("202\torange\tbbb3");
    w.println("302\tstrawberry\tccc3");

    w.close();
    Util.deleteFile(cluster, INPUT_FILE);
    Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);
  }

  @Before
  public void setUp() throws Exception {
    createFile();
    FileSystem fs = FileSystem.getLocal(new Configuration());
    Path localOut = new Path("local-out");
    if (fs.exists(localOut)) {
      fs.delete(localOut, true);
    }
  }

  @After
  public void tearDown() throws Exception {
    new File(INPUT_FILE).delete();
    Util.deleteFile(cluster, INPUT_FILE);
  }

  @AfterClass
  public static void shutdown() {
    cluster.shutDown();
  }

  enum Mode {
    local, cluster
  }

  @Test
  public void testMultiStorage() throws IOException {
    final String LOAD = "A = LOAD '" + INPUT_FILE + "' as (id, name, n);";
    final String MULTI_STORE_CLUSTER = "STORE A INTO 'mr-out' USING "
        + "org.apache.pig.piggybank.storage.MultiStorage('mr-out', '1');";
    final String MULTI_STORE_LOCAL = "STORE A INTO 'local-out' USING "
        + "org.apache.pig.piggybank.storage.MultiStorage('local-out', '1');";

    System.out.print("Testing in LOCAL mode: ...");
    testMultiStorage(Mode.local, "local-out", LOAD, MULTI_STORE_LOCAL);
    System.out.println("Succeeded!");

    System.out.print("Testing in CLUSTER mode: ...");
    testMultiStorage( Mode.cluster, "mr-out", LOAD, MULTI_STORE_CLUSTER);
    System.out.println("Succeeded!");
  }

  @Test
  public void testOutputStats() throws IOException {
    FileSystem fs = cluster.getFileSystem();

    pigServer.setBatchOn();
    pigServer.registerQuery("A = LOAD '" + INPUT_FILE + "' as (id, name, n);");
    pigServer.registerQuery("B = FILTER A BY name == 'apple';");
    pigServer.registerQuery("STORE A INTO 'out1' USING org.apache.pig.piggybank.storage.MultiStorage('out1', '1');"); //153 bytes
    pigServer.registerQuery("STORE B INTO 'out2' USING org.apache.pig.piggybank.storage.MultiStorage('out2', '1');"); // 45 bytes

    ExecJob job = pigServer.executeBatch().get(0);

    PigStats stats = job.getStatistics();
    PigStats.JobGraph jobGraph = stats.getJobGraph();
    JobStats jobStats = (JobStats) jobGraph.getSinks().get(0);
    Map<String, Long> multiStoreCounters = jobStats.getMultiStoreCounters();
    List<OutputStats> outputStats = SimplePigStats.get().getOutputStats();
    OutputStats outputStats1 = "out1".equals(outputStats.get(0).getName()) ? outputStats.get(0) : outputStats.get(1);
    OutputStats outputStats2 = "out2".equals(outputStats.get(0).getName()) ? outputStats.get(0) : outputStats.get(1);

    assertEquals(153 + 45, stats.getBytesWritten());
    assertEquals(2, outputStats.size()); // 2 split conditions
    assertEquals(153, outputStats1.getBytes());
    assertEquals(45, outputStats2.getBytes());
    assertEquals(9, outputStats1.getRecords());
    assertEquals(3, outputStats2.getRecords());
    assertEquals(3L, multiStoreCounters.get("Output records in _1_out2").longValue());
    assertEquals(9L, multiStoreCounters.get("Output records in _0_out1").longValue());

    fs.delete(new Path("out1"), true);
    fs.delete(new Path("out2"), true);
  }

    /**
   * The actual method that run the test in local or cluster mode.
   */
  private void testMultiStorage( Mode mode, String outPath,
      String... queries) throws IOException {
    PigServer pigServer = (Mode.local == mode) ? this.pigServerLocal : this.pigServer;
    pigServer.setBatchOn();
    for (String query : queries) {
      pigServer.registerQuery(query);
    }
    pigServer.executeBatch();
    verifyResults(mode, outPath);
  }

  /**
   * Test if records are split into directories corresponding to split field
   * values
   */
  private void verifyResults(Mode mode, String outPath) throws IOException {
    FileSystem fs = (Mode.local == mode ? FileSystem
        .getLocal(new Configuration()) : cluster.getFileSystem());
    Path output = new Path(outPath);
    assertTrue("Output dir does not exists!", fs.exists(output)
        && fs.getFileStatus(output).isDir());

    Path[] paths = FileUtil.stat2Paths(fs.listStatus(output, hiddenPathFilter));
    assertTrue("Split field dirs not found!", paths != null);

    for (Path path : paths) {
      String splitField = path.getName();
      Path[] files = FileUtil.stat2Paths(fs.listStatus(path, hiddenPathFilter));
      assertTrue("No files found for path: " + path.toUri().getPath(),
          files != null);
      for (Path filePath : files) {
        assertTrue("This shouldn't be a directory", fs.isFile(filePath));
        BufferedReader reader = new BufferedReader(new InputStreamReader(fs
                .open(filePath)));
        String line = "";
        int count = 0;
        while ((line = reader.readLine()) != null) {
          String[] fields = line.split("\\t");
          assertEquals(fields.length, 3);
          assertEquals("Unexpected field value in the output record",
                splitField, fields[1]);
          count++;
          System.out.println("field: " + fields[1]);
        }
        reader.close();
        assertEquals(count, 3);
      }
    }
  }
}