test/org/apache/pig/backend/hadoop/accumulo/TestAccumuloPigCluster.java - pig - Git at Google

 /**
  * 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.backend.hadoop.accumulo;

 import java.io.File;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.TimeUnit;

 import org.apache.accumulo.core.client.BatchWriter;
 import org.apache.accumulo.core.client.BatchWriterConfig;
 import org.apache.accumulo.core.client.Connector;
 import org.apache.accumulo.core.client.ZooKeeperInstance;
 import org.apache.accumulo.core.client.admin.TableOperations;
 import org.apache.accumulo.core.client.security.tokens.PasswordToken;
 import org.apache.accumulo.core.data.Mutation;
 import org.apache.accumulo.minicluster.MiniAccumuloCluster;
 import org.apache.accumulo.minicluster.MiniAccumuloConfig;
 import org.apache.commons.io.FileUtils;
 import org.apache.log4j.Logger;
 import org.apache.pig.PigServer;
 import org.apache.pig.data.Tuple;
 import org.apache.pig.test.Util;
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;

 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.io.Files;

 public class TestAccumuloPigCluster {

     @SuppressWarnings("unchecked")
     private static final List<ImmutableMap<String, String>> AIRPORTS = Lists
             .newArrayList(
                     ImmutableMap.of("code", "SJC", "name", "San Jose"),
                     ImmutableMap.of("code", "SFO", "name", "San Francisco"),
                     ImmutableMap.of("code", "MDO", "name", "Orlando"),
                     ImmutableMap.of("code", "MDW", "name", "Chicago-Midway"),
                     ImmutableMap.of("code", "JFK", "name", "JFK International"),
                     ImmutableMap.of("code", "BWI", "name",
                             "Baltimore-Washington"));

     @SuppressWarnings("unchecked")
     private static final List<ImmutableMap<String, String>> flightData = Lists
             .newArrayList(
                     ImmutableMap.of("origin", "BWI", "destination", "SFO"),
                     ImmutableMap.of("origin", "BWI", "destination", "SJC"),
                     ImmutableMap.of("origin", "MDW", "destination", "MDO"),
                     ImmutableMap.of("origin", "MDO", "destination", "SJC"),
                     ImmutableMap.of("origin", "SJC", "destination", "JFK"),
                     ImmutableMap.of("origin", "JFK", "destination", "MDW"));

     private static final Logger log = Logger
             .getLogger(TestAccumuloPigCluster.class);
     private static final File tmpdir = Files.createTempDir();
     private static MiniAccumuloCluster accumuloCluster;
     private PigServer pig;

     @BeforeClass
     public static void setupClusters() throws Exception {
         MiniAccumuloConfig macConf = new MiniAccumuloConfig(tmpdir, "password");
         macConf.setNumTservers(1);

         accumuloCluster = new MiniAccumuloCluster(macConf);
         accumuloCluster.start();
     }

     @Before
     public void beforeTest() throws Exception {
         pig = new PigServer(Util.getLocalTestMode());
     }

     @AfterClass
     public static void stopClusters() throws Exception {
         accumuloCluster.stop();
         Thread.sleep(5000);
         FileUtils.deleteDirectory(tmpdir);
     }

     private void loadTestData() throws Exception {
         ZooKeeperInstance inst = new ZooKeeperInstance(
                 accumuloCluster.getInstanceName(),
                 accumuloCluster.getZooKeepers());
         Connector c = inst.getConnector("root", new PasswordToken("password"));

         TableOperations tops = c.tableOperations();
         if (!tops.exists("airports")) {
             tops.create("airports");
         }

         if (!tops.exists("flights")) {
             tops.create("flights");
         }

         BatchWriterConfig config = new BatchWriterConfig();
         config.setMaxWriteThreads(1);
         config.setMaxLatency(100000l, TimeUnit.MILLISECONDS);
         config.setMaxMemory(10000l);

         BatchWriter bw = c.createBatchWriter("airports", config);
         try {
             int i = 1;
             for (Map<String, String> record : AIRPORTS) {
                 Mutation m = new Mutation(Integer.toString(i));

                 for (Entry<String, String> entry : record.entrySet()) {
                     m.put(entry.getKey(), "", entry.getValue());
                 }

                 bw.addMutation(m);
                 i++;
             }
         } finally {
             if (null != bw) {
                 bw.close();
             }
         }

         bw = c.createBatchWriter("flights", config);
         try {
             int i = 1;
             for (Map<String, String> record : flightData) {
                 Mutation m = new Mutation(Integer.toString(i));

                 for (Entry<String, String> entry : record.entrySet()) {
                     m.put(entry.getKey(), "", entry.getValue());
                 }

                 bw.addMutation(m);
                 i++;
             }
         } finally {
             if (null != bw) {
                 bw.close();
             }
         }
     }

     @Test
     public void test() throws Exception {
         loadTestData();

         final String loadFlights = "flights = LOAD 'accumulo://flights?instance="
                 + accumuloCluster.getInstanceName()
                 + "&user=root&password=password&zookeepers="
                 + accumuloCluster.getZooKeepers()
                 + "' using org.apache.pig.backend.hadoop.accumulo.AccumuloStorage()"
                 + " as (rowKey:chararray, column_map:map[]);";

         final String loadAirports = "airports = LOAD 'accumulo://airports?instance="
                 + accumuloCluster.getInstanceName()
                 + "&user=root&password=password&zookeepers="
                 + accumuloCluster.getZooKeepers()
                 + "' using org.apache.pig.backend.hadoop.accumulo.AccumuloStorage()"
                 + " as (rowKey:chararray, column_map:map[]);";

         final String joinQuery = "joined = JOIN flights BY column_map#'origin', airports BY column_map#'code';";

         pig.registerQuery(loadFlights);
         pig.registerQuery(loadAirports);
         pig.registerQuery(joinQuery);

         Iterator<Tuple> it = pig.openIterator("joined");

         int i = 0;
         while (it.hasNext()) {
             Tuple t = it.next();

             // id and map for each dataset we joined
             Assert.assertEquals(4, t.size());

             Object o = t.get(1);

             Map<String, String> airport = null, flight = null;

             Assert.assertTrue(Map.class.isAssignableFrom(o.getClass()));
             @SuppressWarnings("unchecked")
             Map<String, String> data1 = (Map<String, String>) o;
             Assert.assertTrue(!data1.isEmpty());

             if (data1.containsKey("origin")) {
                 flight = data1;
             } else if (data1.containsKey("code")) {
                 airport = data1;
             } else {
                 Assert.fail("Received map which did not contain an expected key");
             }

             o = t.get(3);

             Assert.assertTrue(Map.class.isAssignableFrom(o.getClass()));
             @SuppressWarnings("unchecked")
             Map<String, String> data2 = (Map<String, String>) o;
             Assert.assertTrue(!data2.isEmpty());

             if (null == flight && data2.containsKey("origin")) {
                 flight = data2;
             } else if (null == airport && data2.containsKey("code")) {
                 airport = data2;
             } else {
                 Assert.fail("Received map which did not contain an expected key");
             }

             Assert.assertTrue(null != airport && null != flight);

             Assert.assertEquals(airport.get("code"), flight.get("origin"));

             i++;
         }

         Assert.assertEquals(6, i);
     }

 }
	/**
	* 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.backend.hadoop.accumulo;

	import java.io.File;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import java.util.Map.Entry;
	import java.util.concurrent.TimeUnit;

	import org.apache.accumulo.core.client.BatchWriter;
	import org.apache.accumulo.core.client.BatchWriterConfig;
	import org.apache.accumulo.core.client.Connector;
	import org.apache.accumulo.core.client.ZooKeeperInstance;
	import org.apache.accumulo.core.client.admin.TableOperations;
	import org.apache.accumulo.core.client.security.tokens.PasswordToken;
	import org.apache.accumulo.core.data.Mutation;
	import org.apache.accumulo.minicluster.MiniAccumuloCluster;
	import org.apache.accumulo.minicluster.MiniAccumuloConfig;
	import org.apache.commons.io.FileUtils;
	import org.apache.log4j.Logger;
	import org.apache.pig.PigServer;
	import org.apache.pig.data.Tuple;
	import org.apache.pig.test.Util;
	import org.junit.AfterClass;
	import org.junit.Assert;
	import org.junit.Before;
	import org.junit.BeforeClass;
	import org.junit.Test;

	import com.google.common.collect.ImmutableMap;
	import com.google.common.collect.Lists;
	import com.google.common.io.Files;

	public class TestAccumuloPigCluster {

	@SuppressWarnings("unchecked")
	private static final List<ImmutableMap<String, String>> AIRPORTS = Lists
	.newArrayList(
	ImmutableMap.of("code", "SJC", "name", "San Jose"),
	ImmutableMap.of("code", "SFO", "name", "San Francisco"),
	ImmutableMap.of("code", "MDO", "name", "Orlando"),
	ImmutableMap.of("code", "MDW", "name", "Chicago-Midway"),
	ImmutableMap.of("code", "JFK", "name", "JFK International"),
	ImmutableMap.of("code", "BWI", "name",
	"Baltimore-Washington"));

	@SuppressWarnings("unchecked")
	private static final List<ImmutableMap<String, String>> flightData = Lists
	.newArrayList(
	ImmutableMap.of("origin", "BWI", "destination", "SFO"),
	ImmutableMap.of("origin", "BWI", "destination", "SJC"),
	ImmutableMap.of("origin", "MDW", "destination", "MDO"),
	ImmutableMap.of("origin", "MDO", "destination", "SJC"),
	ImmutableMap.of("origin", "SJC", "destination", "JFK"),
	ImmutableMap.of("origin", "JFK", "destination", "MDW"));

	private static final Logger log = Logger
	.getLogger(TestAccumuloPigCluster.class);
	private static final File tmpdir = Files.createTempDir();
	private static MiniAccumuloCluster accumuloCluster;
	private PigServer pig;

	@BeforeClass
	public static void setupClusters() throws Exception {
	MiniAccumuloConfig macConf = new MiniAccumuloConfig(tmpdir, "password");
	macConf.setNumTservers(1);

	accumuloCluster = new MiniAccumuloCluster(macConf);
	accumuloCluster.start();
	}

	@Before
	public void beforeTest() throws Exception {
	pig = new PigServer(Util.getLocalTestMode());
	}

	@AfterClass
	public static void stopClusters() throws Exception {
	accumuloCluster.stop();
	Thread.sleep(5000);
	FileUtils.deleteDirectory(tmpdir);
	}

	private void loadTestData() throws Exception {
	ZooKeeperInstance inst = new ZooKeeperInstance(
	accumuloCluster.getInstanceName(),
	accumuloCluster.getZooKeepers());
	Connector c = inst.getConnector("root", new PasswordToken("password"));

	TableOperations tops = c.tableOperations();
	if (!tops.exists("airports")) {
	tops.create("airports");
	}

	if (!tops.exists("flights")) {
	tops.create("flights");
	}

	BatchWriterConfig config = new BatchWriterConfig();
	config.setMaxWriteThreads(1);
	config.setMaxLatency(100000l, TimeUnit.MILLISECONDS);
	config.setMaxMemory(10000l);

	BatchWriter bw = c.createBatchWriter("airports", config);
	try {
	int i = 1;
	for (Map<String, String> record : AIRPORTS) {
	Mutation m = new Mutation(Integer.toString(i));

	for (Entry<String, String> entry : record.entrySet()) {
	m.put(entry.getKey(), "", entry.getValue());
	}

	bw.addMutation(m);
	i++;
	}
	} finally {
	if (null != bw) {
	bw.close();
	}
	}

	bw = c.createBatchWriter("flights", config);
	try {
	int i = 1;
	for (Map<String, String> record : flightData) {
	Mutation m = new Mutation(Integer.toString(i));

	for (Entry<String, String> entry : record.entrySet()) {
	m.put(entry.getKey(), "", entry.getValue());
	}

	bw.addMutation(m);
	i++;
	}
	} finally {
	if (null != bw) {
	bw.close();
	}
	}
	}

	@Test
	public void test() throws Exception {
	loadTestData();

	final String loadFlights = "flights = LOAD 'accumulo://flights?instance="
	+ accumuloCluster.getInstanceName()
	+ "&user=root&password=password&zookeepers="
	+ accumuloCluster.getZooKeepers()
	+ "' using org.apache.pig.backend.hadoop.accumulo.AccumuloStorage()"
	+ " as (rowKey:chararray, column_map:map[]);";

	final String loadAirports = "airports = LOAD 'accumulo://airports?instance="
	+ accumuloCluster.getInstanceName()
	+ "&user=root&password=password&zookeepers="
	+ accumuloCluster.getZooKeepers()
	+ "' using org.apache.pig.backend.hadoop.accumulo.AccumuloStorage()"
	+ " as (rowKey:chararray, column_map:map[]);";

	final String joinQuery = "joined = JOIN flights BY column_map#'origin', airports BY column_map#'code';";

	pig.registerQuery(loadFlights);
	pig.registerQuery(loadAirports);
	pig.registerQuery(joinQuery);

	Iterator<Tuple> it = pig.openIterator("joined");

	int i = 0;
	while (it.hasNext()) {
	Tuple t = it.next();

	// id and map for each dataset we joined
	Assert.assertEquals(4, t.size());

	Object o = t.get(1);

	Map<String, String> airport = null, flight = null;

	Assert.assertTrue(Map.class.isAssignableFrom(o.getClass()));
	@SuppressWarnings("unchecked")
	Map<String, String> data1 = (Map<String, String>) o;
	Assert.assertTrue(!data1.isEmpty());

	if (data1.containsKey("origin")) {
	flight = data1;
	} else if (data1.containsKey("code")) {
	airport = data1;
	} else {
	Assert.fail("Received map which did not contain an expected key");
	}

	o = t.get(3);

	Assert.assertTrue(Map.class.isAssignableFrom(o.getClass()));
	@SuppressWarnings("unchecked")
	Map<String, String> data2 = (Map<String, String>) o;
	Assert.assertTrue(!data2.isEmpty());

	if (null == flight && data2.containsKey("origin")) {
	flight = data2;
	} else if (null == airport && data2.containsKey("code")) {
	airport = data2;
	} else {
	Assert.fail("Received map which did not contain an expected key");
	}

	Assert.assertTrue(null != airport && null != flight);

	Assert.assertEquals(airport.get("code"), flight.get("origin"));

	i++;
	}

	Assert.assertEquals(6, i);
	}

	}