/**
* 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);
}
}