package com.twitter.elephantbird.pig.load;

import com.google.common.collect.Maps;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.input.TextInputFormat;
import org.apache.pig.PigException;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.PigContext;
import org.json.simple.JSONObject;
import org.json.simple.parser.JSONParser;
import org.json.simple.parser.ParseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Map;

/**
 * Decodes each line as JSON passes the resulting map of values
 * to Pig as a single-element tuple.
 */
public class JsonLoader extends LzoBaseLoadFunc {
  private static final Logger LOG = LoggerFactory.getLogger(JsonLoader.class);
  private static final TupleFactory tupleFactory = TupleFactory.getInstance();

  private final JSONParser jsonParser = new JSONParser();

  private enum JsonLoaderCounters {
    LinesRead,
    LinesJsonDecoded,
    LinesParseError,
    LinesParseErrorBadNumber
  }

  private String inputFormatClassName;

  public JsonLoader() {
    this(TextInputFormat.class.getName());
  }

  public JsonLoader(String inputFormatClassName) {
    this.inputFormatClassName = inputFormatClassName;
  }

  /**
   * Return every non-null line as a single-element tuple to Pig.
   */
  @Override
  public Tuple getNext() throws IOException {
    if (reader == null) {
      return null;
    }
    try {
      while (reader.nextKeyValue()) {
        Text value = (Text) reader.getCurrentValue();
        incrCounter(JsonLoaderCounters.LinesRead, 1L);
        Tuple t = parseStringToTuple(value.toString());
        if (t != null) {
          incrCounter(JsonLoaderCounters.LinesJsonDecoded, 1L);
          return t;
        }
      }
      return null;

    } catch (InterruptedException e) {
      int errCode = 6018;
      String errMsg = "Error while reading input";
      throw new ExecException(errMsg, errCode,
          PigException.REMOTE_ENVIRONMENT, e);
    }
  }

  @Override
  public InputFormat getInputFormat() throws IOException {
    try {
      return (FileInputFormat) PigContext.resolveClassName(inputFormatClassName).newInstance();
    } catch (InstantiationException e) {
      throw new IOException("Failed creating input format " + inputFormatClassName, e);
    } catch (IllegalAccessException e) {
      throw new IOException("Failed creating input format " + inputFormatClassName, e);
    }
  }

  protected Tuple parseStringToTuple(String line) {
    try {
      Map<String, String> values = Maps.newHashMap();
      JSONObject jsonObj = (JSONObject) jsonParser.parse(line);
      if (jsonObj != null) {
        for (Object key : jsonObj.keySet()) {
          Object value = jsonObj.get(key);
          values.put(key.toString(), value != null ? value.toString() : null);
        }
      } else {
        // JSONParser#parse(String) may return a null reference, e.g. when
        // the input parameter is the string "null".  A single line with
        // "null" is not valid JSON though.
        LOG.warn("Could not json-decode string: " + line);
        incrCounter(JsonLoaderCounters.LinesParseError, 1L);
        return null;
      }
      return tupleFactory.newTuple(values);
    } catch (ParseException e) {
      LOG.warn("Could not json-decode string: " + line, e);
      incrCounter(JsonLoaderCounters.LinesParseError, 1L);
      return null;
    } catch (NumberFormatException e) {
      LOG.warn("Very big number exceeds the scale of long: " + line, e);
      incrCounter(JsonLoaderCounters.LinesParseErrorBadNumber, 1L);
      return null;
    } catch (ClassCastException e) {
      LOG.warn("Could not convert to Json Object: " + line, e);
      incrCounter(JsonLoaderCounters.LinesParseError, 1L);
      return null;
    }
  }

}

@@ -371,32 +371,153 @@
 .ExecException;%0A
+import org.apache.pig.data.BagFactory;%0Aimport org.apache.pig.data.DataBag;%0Aimport org.apache.pig.data.DefaultBagFactory;%0A
 import org.apach
@@ -534,16 +534,16 @@
 .Tuple;%0A
-
 import o
@@ -610,24 +610,58 @@
 PigContext;%0A
+import org.json.simple.JSONArray;%0A
 import org.j
@@ -1196,16 +1196,96 @@
 tance();
+%0A  private static final BagFactory bagFactory = DefaultBagFactory.getInstance();
 %0A%0A  priv
@@ -2967,62 +2967,8 @@
 y %7B%0A
-      Map%3CString, String%3E values = Maps.newHashMap();%0A
     
@@ -3068,177 +3068,56 @@
     
-for (Object key : jsonObj.keySet()) %7B%0A          Object value = jsonObj.get(key);%0A          values.put(key.toString(), value != null ? value.toString() : null);%0A        %7D
+return tupleFactory.newTuple(walkJson(jsonObj));
 %0A   
@@ -3473,52 +3473,8 @@
   %7D%0A
-      return tupleFactory.newTuple(values);%0A
     
@@ -4030,10 +4030,744 @@
  %7D%0A  %7D%0A%0A
+  private Object wrap(Object value) %7B%0A    %0A    if (value instanceof JSONObject) %7B%0A      return walkJson((JSONObject) value);%0A    %7D  else if (value instanceof JSONArray) %7B%0A      %0A      JSONArray a = (JSONArray) value;%0A      DataBag mapValue = bagFactory.newDefaultBag();%0A      for (int i=0; i%3Ca.size(); i++) %7B%0A        Tuple t = tupleFactory.newTuple(wrap(a.get(i)));%0A        mapValue.add(t);%0A      %7D%0A      return mapValue;%0A      %0A    %7D else %7B%0A      return value != null ? value.toString() : null;%0A    %7D%0A  %7D%0A%0A  private Map%3CString,Object%3E walkJson(JSONObject jsonObj) %7B%0A    Map%3CString,Object%3E v = Maps.newHashMap();%0A    for (Object key: jsonObj.keySet()) %7B%0A      v.put(key.toString(), wrap(jsonObj.get(key)));%0A    %7D%0A    return v;%0A  %7D%0A  %0A
 %7D%0A

Load nested json structures: objects and arrays. Objects are loaded as nested Maps. Arrays are loaded as Bags.