/*
 * Copyright 2015 SURFsara B.V.
 *
 * Licensed 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 nl.surfsara.streamcorpus;

import com.twitter.elephantbird.mapreduce.io.ThriftWritable;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TIOStreamTransport;
import streamcorpus.StreamItem;

import java.io.IOException;

public class StreamCorpusRecordReader extends RecordReader<Text, ThriftWritable<StreamItem>> {

    private FileSplit fileSplit;
    private Configuration conf;

    private TProtocol tProtocol;

    private Text key;
    private ThriftWritable<StreamItem> value;
    private StreamItem streamItem;

    @Override
    public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
        fileSplit = (FileSplit) inputSplit;
        conf = taskAttemptContext.getConfiguration();

        key = new Text();
        value = new ThriftWritable<>();
        streamItem = new StreamItem();

        final Path file = fileSplit.getPath();
        FileSystem fs = file.getFileSystem(conf);
        TIOStreamTransport in = new TIOStreamTransport(fs.open(file));
        tProtocol = new TBinaryProtocol.Factory().getProtocol(in);
    }

    @Override
    public boolean nextKeyValue() throws IOException, InterruptedException {
        key.set(fileSplit.getPath().toString());

        try {
            streamItem.read(tProtocol);
            value.set(streamItem);
        } catch (TException e) {
            throw new IOException();
        }

        if (streamItem == null) {
            return false;
        }

        return true;
    }

    @Override
    public Text getCurrentKey() throws IOException, InterruptedException {
        return key;
    }

    @Override
    public ThriftWritable<StreamItem> getCurrentValue() throws IOException, InterruptedException {
        return value;
    }

    @Override
    public float getProgress() throws IOException, InterruptedException {
        return 0f;
    }

    @Override
    public void close() throws IOException {
        tProtocol.getTransport().close();
    }
}

@@ -783,48 +783,8 @@
 am;%0A
-import org.apache.hadoop.fs.FileSystem;%0A
 impo
@@ -1242,16 +1242,72 @@
 nsport;%0A
+import org.apache.thrift.transport.TTransportException;%0A
 import s
@@ -1478,21 +1478,17 @@
 leSplit 
-fileS
+s
 plit;%0A  
@@ -1541,26 +1541,48 @@
 otocol t
-Protocol;%0A
+p;%0A%0A    private StreamItem item;
 %0A    pri
@@ -1645,43 +1645,8 @@
 lue;
-%0A    private StreamItem streamItem;
 %0A%0A  
@@ -1797,21 +1797,17 @@
         
-fileS
+s
 plit = (
@@ -1887,115 +1887,8 @@
 ();%0A
-%0A        key = new Text();%0A        value = new ThriftWritable%3C%3E();%0A        streamItem = new StreamItem();%0A%0A
     
@@ -1909,21 +1909,17 @@
  file = 
-fileS
+s
 plit.get
@@ -1939,20 +1939,27 @@
    F
-ileSystem fs
+SDataInputStream in
  = f
@@ -1985,67 +1985,8 @@
 onf)
-;%0A        TIOStreamTransport in = new TIOStreamTransport(fs
 .ope
@@ -1992,17 +1992,16 @@
 en(file)
-)
 ;%0A      
@@ -2003,24 +2003,17 @@
        t
-Protocol
+p
  = new T
@@ -2053,11 +2053,35 @@
 col(
+new TIOStreamTransport(
 in)
+)
 ;%0A  
@@ -2168,32 +2168,99 @@
 ptedException %7B%0A
+        if (key == null) %7B%0A            key = new Text();%0A        %7D%0A
         key.set(
@@ -2263,13 +2263,9 @@
 set(
-fileS
+s
 plit
@@ -2288,25 +2288,24 @@
 ing());%0A
-%0A
         
 try %7B%0A  
@@ -2300,51 +2300,196 @@
     
-try %7B%0A            streamItem.read(tProtocol
+if (item == null) %7B%0A            item = new StreamItem();%0A        %7D%0A        if (value == null) %7B%0A            value = new ThriftWritable%3C%3E();%0A        %7D%0A        try %7B%0A            item.read(tp
 );%0A 
@@ -2509,23 +2509,17 @@
 lue.set(
-streamI
+i
 tem);%0A  
@@ -2565,32 +2565,121 @@
     
-throw new IOException();
+if (e instanceof TTransportException && ((TTransportException) e).getType() == TTransportException.END_OF_FILE) %7B
 %0A   
@@ -2683,27 +2683,24 @@
         
-%7D%0A%0A
         
 if (stre
@@ -2695,59 +2695,134 @@
     
-if (streamItem == null) %7B%0A            return false;
+return false;%0A            %7D else %7B%0A                e.printStackTrace();%0A                throw new IOException();%0A            %7D
 %0A   
@@ -2824,25 +2824,24 @@
 %7D%0A        %7D%0A
-%0A
         retu
@@ -3295,25 +3295,77 @@
         
-tProtocol
+if (tp != null && tp.getTransport() != null) %7B%0A            tp
 .getTran
@@ -3377,20 +3377,30 @@
 ().close();%0A
+        %7D%0A
     %7D%0A%7D%0A

Handle EOF correctly for Thrift stream