it.crs4.seal.common.TestContext.java Source code

Introduction

Here is the source code for it.crs4.seal.common.TestContext.java

Source

// Copyright (C) 2011-2012 CRS4.
//
// This file is part of Seal.
//
// Seal is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option)
// any later version.
//
// Seal is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
// for more details.
//
// You should have received a copy of the GNU General Public License along
// with Seal.  If not, see <http://www.gnu.org/licenses/>.

package it.crs4.seal.common;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;
import java.util.List;
import java.util.Iterator;

import org.apache.hadoop.io.Writable;

public class TestContext<K, V> implements IMRContext<K, V> {
    private static class CounterStore {
        private HashMap<String, HashMap<String, Long>> counters;

        public CounterStore() {
            counters = new HashMap<String, HashMap<String, Long>>(5);
        }

        public long getValue(String group, String name) {
            HashMap<String, Long> map = counters.get(group);
            if (map != null) {
                Long value = map.get(name);
                if (value != null)
                    return value;
            }
            return 0;
        }

        public void setValue(String group, String name, long value) {
            HashMap<String, Long> map = counters.get(group);
            if (map == null) {
                map = new HashMap<String, Long>();
                counters.put(group, map);
            }
            map.put(name, value);
        }
    }

    public static class Tuple<K, V> {
        public K key;
        public V value;

        public Tuple(K k, V v) {
            key = k;
            value = v;
        }

        public K getKey() {
            return key;
        }

        public V getValue() {
            return value;
        }
    }

    protected ArrayList<Tuple<K, V>> output;

    protected int progressCalled;
    protected String lastStatus;
    protected CounterStore counters;

    public TestContext() {
        output = new ArrayList<Tuple<K, V>>(30);
        progressCalled = 0;
        counters = new CounterStore();
    }

    public void progress() {
        progressCalled += 1;
    }

    public boolean getProgressCalled() {
        return progressCalled > 0;
    }

    public int getNumProgressCalls() {
        return progressCalled;
    }

    public void setStatus(String msg) {
        lastStatus = msg;
    }

    public String getLastStatus() {
        return lastStatus;
    }

    @SuppressWarnings("unchecked")
    public void write(K key, V value) throws java.io.IOException, InterruptedException {
        // If possible, duplicate the objects we store to sever dependencies to the mapper or reducer objects,
        // like the real Hadoop context does when values are trasferred between map and reduce phases.
        if (key instanceof Writable && value instanceof Writable)
            output.add(new Tuple<K, V>((K) duplicateWritable(key), (V) duplicateWritable(value))); // unchecked casts that generate warnings
        else
            output.add(new Tuple<K, V>(key, value));
    }

    public Set<K> getKeys() {
        Set<K> set = new HashSet<K>();
        for (Tuple<K, V> pair : output)
            set.add(pair.key);

        return set;
    }

    public List<V> getValuesForKey(K key) {
        List<V> list = new ArrayList<V>();
        for (Tuple<K, V> pair : output) {
            if (key == pair.key || key != null && key.equals(pair.key))
                list.add(pair.value);
        }

        return list;
    }

    public List<V> getAllValues() {
        List<V> list = new ArrayList<V>();
        for (Tuple<K, V> pair : output)
            list.add(pair.value);

        return list;
    }

    public Iterator<Tuple<K, V>> iterator() {
        return output.iterator();
    }

    public int getNumWrites() {
        return output.size();
    }

    public void increment(Enum<?> counter, long value) {
        increment(counter.getClass().getName(), counter.name(), value);
    }

    public void increment(String groupName, String counterName, long value) {
        long currentValue = counters.getValue(groupName, counterName);
        counters.setValue(groupName, counterName, currentValue + value);
    }

    public long getCounterValue(String groupName, String counterName) {
        return counters.getValue(groupName, counterName);
    }

    /**
     * Duplicate Writable object by by serializing and then unserializing it.
     */
    private Object duplicateWritable(Object oldItem) throws java.io.IOException {
        if (oldItem == null)
            return null;

        Writable w = (Writable) oldItem;
        // duplicate the key by serializing and then unserializing it
        ByteArrayOutputStream obytes = new ByteArrayOutputStream();
        DataOutputStream ostream = new DataOutputStream(obytes);

        w.write(ostream);
        ostream.close();

        Object newItem;
        try {
            newItem = oldItem.getClass().newInstance();
        } catch (Exception e) {
            throw new RuntimeException("error instantiating duplicate key or value.  Message: " + e.getMessage());
        }

        ByteArrayInputStream ibytes = new ByteArrayInputStream(obytes.toByteArray());
        DataInputStream istream = new DataInputStream(ibytes);
        ((Writable) newItem).readFields(istream);
        return newItem;
    }
}