Graph generator with HDFS support

build.sbt

@@ -24,6 +24,8 @@ libraryDependencies += "org.mockito" % "mockito-all" % "1.8.5" % "test" withSour
 
 libraryDependencies += "org.scala-tools.testing" % "specs_2.8.1" % "1.6.6" % "test" withSources()
 
+libraryDependencies += "org.apache.hadoop" % "hadoop-core" % "0.20.2"
+
 publishMavenStyle := true
 
 publishTo <<= version { (v: String) =>

cassovary.iml

@@ -5,6 +5,8 @@
       <configuration>
         <option name="compilerLibraryLevel" value="Project" />
         <option name="compilerLibraryName" value="buildScala" />
+        <option name="maximumHeapSize" value="4096" />
+        <option name="vmOptions" value="-Xss256m -server" />
       </configuration>
     </facet>
   </component>

src/main/java/com/twitter/pers/graph_generator/EdgeListOutput.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2012 Twitter, Inc.
+ *
+ * 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 com.twitter.pers.graph_generator;
+
+import java.io.BufferedOutputStream;
+import java.io.DataOutputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+
+/**
+ * Efficient parallel output of edges into a binary edge list format.
+ * @author Aapo Kyrola, [email protected], [email protected]
+ */
+public class EdgeListOutput implements GraphOutput {
+
+    private String fileNamePrefix;
+
+    static int partSeq = 0;
+
+    public EdgeListOutput(String fileNamePrefix) {
+        this.fileNamePrefix = fileNamePrefix;
+    }
+
+    @Override
+    public void addEdges(int[] from, int[] to)  {
+        try {
+            DataOutputStream dos = partitionOut.get();
+            int n = from.length;
+            for(int i=0; i<n; i++) {
+                dos.writeInt(Integer.reverseBytes(from[i]));
+                dos.writeInt(Integer.reverseBytes(to[i]));
+            }
+        } catch (IOException ioe) {
+            throw new RuntimeException(ioe);
+        }
+    }
+
+    public void finishUp() {
+        try {
+            partitionOut.get().close();
+        } catch (IOException ioe) {
+            throw new RuntimeException(ioe);
+        }
+    }
+
+    /* Each thread will have a local partition */
+    private ThreadLocal<DataOutputStream> partitionOut = new ThreadLocal<DataOutputStream>() {
+        @Override
+        protected DataOutputStream initialValue() {
+            try {
+                int thisPartId;
+                synchronized (this) {
+                    thisPartId = partSeq++;
+                }
+
+                String fileName = fileNamePrefix + "-part" + thisPartId;
+                return new DataOutputStream(new BufferedOutputStream(new FileOutputStream(fileName)));
+            } catch (Exception err) {
+                err.printStackTrace();
+                throw new RuntimeException(err);
+            }
+        }
+    };
+
+}

src/main/java/com/twitter/pers/graph_generator/GraphOutput.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2012 Twitter, Inc.
+ *
+ * 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 com.twitter.pers.graph_generator;
+
+/**
+ * @author Aapo Kyrola, [email protected], [email protected]
+ */
+public interface GraphOutput {
+
+    void addEdges(int[] from, int[] to);
+
+    void finishUp();
+
+}

src/main/java/com/twitter/pers/graph_generator/HDFSEdgeListOutput.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2012 Twitter, Inc.
+ *
+ * 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 com.twitter.pers.graph_generator;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+
+import java.io.IOException;
+
+/**
+ * Outputs edges into HDFS in tab delimited edge list format
+ * @author Aapo Kyrola, [email protected], [email protected]
+ */
+public class HDFSEdgeListOutput implements GraphOutput {
+
+    private String fileNamePrefix;
+
+    static int partSeq = 0;
+
+    public HDFSEdgeListOutput(String fileNamePrefix) {
+        this.fileNamePrefix = fileNamePrefix;
+        System.out.println("Using HDFS output: " + fileNamePrefix);
+    }
+
+    @Override
+    public void addEdges(int[] from, int[] to)  {
+        try {
+            FSDataOutputStream dos = partitionOut.get();
+            int n = from.length;
+            StringBuffer sb = new StringBuffer(from.length * 32);
+            for(int i=0; i<n; i++) {
+                sb.append(from[i]);
+                sb.append("\t");
+                sb.append(to[i]);
+                sb.append("\n");
+            }
+            dos.write(sb.toString().getBytes());
+        } catch (Exception ioe) {
+            throw new RuntimeException(ioe);
+        }
+    }
+
+    public void finishUp() {
+        try {
+            partitionOut.get().close();
+        } catch (IOException ioe) {
+            throw new RuntimeException(ioe);
+        }
+    }
+
+    /* Each thread will have a local partition */
+    private ThreadLocal<FSDataOutputStream> partitionOut = new ThreadLocal<FSDataOutputStream>() {
+        @Override
+        protected FSDataOutputStream initialValue() {
+            try {
+                int thisPartId;
+                synchronized (this) {
+                    thisPartId = partSeq++;
+                }
+
+                String hadoopHome = System.getProperty("HADOOP_HOME");
+                if (hadoopHome == null) hadoopHome = System.getenv("HADOOP_HOME");
+
+                if (hadoopHome == null) {
+                    throw new IllegalArgumentException("You need to specify environment variable or JVM option HADOOP_HOME!");
+                }
+
+                Configuration conf = new Configuration();
+                conf.addResource(new Path(hadoopHome + "/conf/core-site.xml"));
+                conf.addResource(new Path(hadoopHome + "/conf/hdfs-site.xml"));
+
+
+                String fileName = fileNamePrefix + "-part" + thisPartId;
+                FileSystem fs = FileSystem.get(conf);
+                return fs.create(new Path(fileName));
+            } catch (Exception err) {
+                err.printStackTrace();
+                throw new RuntimeException(err);
+            }
+        }
+    };
+
+}