/*
 * 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.archive.hadoop.pig;

import java.io.IOException;

import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;

import org.apache.pig.StoreFunc;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;

import org.archive.hadoop.mapreduce.ZipNumOutputFormat;

/**
 * Very simple/minimal StoreFunc to write {key,value} pairs into IA's
 * "zip num" format, using the ZipNumOutputFormat class.
 */
public class ZipNumStorage extends StoreFunc
{
  RecordWriter writer;

  Text key   = new Text();
  Text value = new Text();
  
  public OutputFormat getOutputFormat() throws IOException
  {
    return new ZipNumOutputFormat();
  }
  
  public void setStoreLocation( String location, Job job ) throws IOException
  {
    FileOutputFormat.setOutputPath( job, new Path(location) );
  }

  public void prepareToWrite( RecordWriter writer) throws IOException  
  {
    this.writer = writer;
  }

  /**
   * Tuples must have 2 (and only 2) chararray fields.  This is due to
   * the requirements of ZipNumOutputFormat's RecordWriter.
   */
  public void putNext( Tuple tuple ) throws IOException
  {
    try
      {
        int size = tuple.size();
        
        if ( size != 2 ) throw new IOException( "Tuple size != 2" );
        
        if ( DataType.findType( tuple.get(0) ) != DataType.CHARARRAY )
          {
            throw new IOException( "Invalid type for tuple 0, not CHARARRAY: " + DataType.findTypeName( DataType.findType( tuple.get(0) ) ) + ":" + tuple.get(0) + ":" + tuple.get(1) );
          }
        
        if ( DataType.findType( tuple.get(1) ) != DataType.CHARARRAY )
          {
            throw new IOException( "Invalid type for tuple 1, not CHARARRAY: " + DataType.findTypeName( DataType.findType( tuple.get(1) ) ) + ":" + tuple.get(0) + ":" + tuple.get(1) );
          }
        
        this.key  .set( (String) tuple.get(0) );
        this.value.set( (String) tuple.get(1) );
        
        this.writer.write( this.key, this.value );
      }
    catch ( InterruptedException ie )
      {
        throw new IOException( ie );
      }
  }
}

@@ -1944,32 +1944,34 @@
 les 
-must
+can
  have 
-2 (and only 2)
+either one or two
  cha
@@ -1981,90 +1981,304 @@
 ray 
-field
+value
 s.
+%0A
   
-This is due to
+ *%0A   * If there is only 1, then we split it on the first space.
 %0A   *
+ If
  the
- 
 re
-quirements of ZipNumOutputFormat's RecordWriter
+ are 2, then we just accept them as-is.%0A   *%0A   * The underlying ZipNumOutputFormat requires that the data%0A   * be split into a (key,value) pair, with both the key and%0A   * value being Strings (Pig chararray)
 .%0A  
@@ -2390,24 +2390,25 @@
 size();%0A
+%0A
         
 %0A       
@@ -2403,77 +2403,145 @@
     
-%0A        if ( size != 2 ) throw new IOException( %22Tuple size != 2%22 );
+if ( size != 1 && size != 2 )%0A          %7B%0A            throw new IOException( %22Invalid tuple size, must be 1 or 2: %22 + size );%0A          %7D
 %0A   
@@ -2839,32 +2839,73 @@
    %0A        if (
+ size == 2 )%0A          %7B%0A            if (
  DataType.findTy
@@ -2949,34 +2949,42 @@
 RAY )%0A          
-%7B%0A
+    %7B%0A    
             thro
@@ -3158,34 +3158,40 @@
 ;%0A          
-%7D%0A
+    %7D%0A  
         
 %0A        thi
@@ -3170,33 +3170,39 @@
     %7D%0A          
-%0A
+  %0A    
         this.key
@@ -3238,24 +3238,28 @@
  );%0A        
+   
+ 
 this.value.s
@@ -3279,32 +3279,231 @@
 tuple.get(1) );%0A
+          %7D%0A        else%0A          %7B%0A            String s%5B%5D = ((String)tuple.get(0)).split( %22 %22, 2 );%0A            %0A            this.key  .set( s%5B0%5D );%0A            this.value.set( s%5B1%5D );%0A          %7D%0A
         %0A       

Allow both 2-value and 1-value tuples. If 2-value, assume they are already in (key,value) form. If 1-value, split on first space into (key,value).