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,63 @@
+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,12 @@
+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,83 @@
+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);
+            }
+        }
+    };
+
+}

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

@@ -0,0 +1,166 @@
+package com.twitter.pers.graph_generator;
+
+import java.util.ArrayList;
+import java.util.Random;
+
+/**
+ * Graph generator based on the R-MAT algorithm
+ * R-MAT: A Recursive Model for Graph Mining
+ * Chakrabarti, Zhan, Faloutsos: http://www.cs.cmu.edu/~christos/PUBLICATIONS/siam04.pdf
+ *
+ * Usage:
+ *    [outputfile-prefix] [num-vertices] [num-edges] probA probB probC probD
+ *
+ * See the paper for the description of parameters probA, probB, probC, probD.
+ *
+ * If the outputfile-prefix starts with hdfs://, the graph is written into HDFS as tab-delimited
+ * partitions.
+ *
+ * Specify the number of threads used by passing a JVM option -Dnum_threads
+ * Each thread will create its own partition.
+ *
+ * Note that the result may (will) contain duplicate edges.
+ */
+public class RMATGraphGenerator {
+
+    private GraphOutput outputter;
+
+    /* Parameters for top-left, top-right, bottom-left, bottom-right probabilities */
+    private double pA, pB, pC, pD;
+    private long numEdges;
+    private int numVertices;
+
+    /**
+     *  From http://pywebgraph.sourceforge.net
+     ## Probability of choosing quadrant A
+     self.probA = 0.45
+
+     ## Probability of choosing quadrant B
+     self.probB = 0.15
+
+     ## Probability of choosing quadrant C
+     self.probC = 0.15
+
+     ## Probability of choosing quadrant D
+     self.probD = 0.25
+     */
+
+
+    public RMATGraphGenerator(GraphOutput outputter, double pA, double pB, double pC, double pD, int nVertices,long nEdges) {
+        this.outputter = outputter;
+        this.pA = pA;
+        this.pB = pB;
+        this.pC = pC;
+        this.pD = pD;
+
+        if (Math.abs(pA + pB + pC + pD - 1.0) > 0.01)
+            throw new IllegalArgumentException("Probabilities do not add up to one!");
+        numVertices = nVertices;
+        numEdges = nEdges;
+    }
+
+    public void execute() {
+        int nThreads = Integer.parseInt(System.getProperty("num_threads", Runtime.getRuntime().availableProcessors() + ""));
+
+        ArrayList<Thread> threads = new ArrayList<Thread>();
+        for(int i=0; i < nThreads; i++) {
+            Thread t = new Thread(new RMATGenerator(numEdges / nThreads));
+            t.start();
+            threads.add(t);
+        }
+
+        /* Wait */
+        try {
+            for(Thread t : threads) t.join();
+        } catch (InterruptedException ie) {
+            throw new RuntimeException(ie);
+        }
+    }
+
+    private class RMATGenerator implements Runnable {
+
+        private long edgesToGenerate;
+
+        private RMATGenerator(long genEdges) {
+            this.edgesToGenerate = genEdges;
+        }
+
+        public void run() {
+            int nEdgesATime = 1000000;
+            long createdEdges = 0;
+
+            Random r = new Random(System.currentTimeMillis() + this.hashCode());
+
+            double cumA = pA;
+            double cumB = cumA + pB;
+            double cumC = cumB + pC;
+            double cumD = 1.0;
+            assert(cumD > cumC);
+
+            while(edgesToGenerate > createdEdges) {
+                int ne = (int) Math.min(edgesToGenerate  - createdEdges, nEdgesATime);
+                int[] fromIds = new int[ne];
+                int[] toIds = new int[ne];
+
+                for(int j=0; j < ne; j++) {
+                    int col_st = 0, col_en = numVertices - 1, row_st = 0, row_en = numVertices - 1;
+                    while (col_st != col_en || row_st != row_en) {
+                        double x = r.nextDouble();
+
+                        if (x < cumA) {
+                            // Top-left
+                            col_en = col_st + (col_en - col_st) / 2;
+                            row_en = row_st + (row_en - row_st) / 2;
+                        } else if (x < cumB) {
+                            // Top-right
+                            col_st = col_en - (col_en - col_st) / 2;
+                            row_en = row_st + (row_en - row_st) / 2;
+
+                        } else if (x < cumC) {
+                            // Bottom-left
+                            col_en = col_st + (col_en - col_st) / 2;
+                            row_st = row_en - (row_en - row_st) / 2;
+                        } else {
+                            // Bottom-right
+                            col_st = col_en - (col_en - col_st) / 2;
+                            row_st = row_en - (row_en - row_st) / 2;
+                        }
+                    }
+                    fromIds[j] = col_st;
+                    toIds[j] = row_st;
+                }
+
+                outputter.addEdges(fromIds,  toIds);
+                createdEdges += ne;
+                System.out.println(Thread.currentThread().getId() + " created " + createdEdges + " edges.");
+            }
+            outputter.finishUp();
+        }
+    }
+
+    public static void main(String[] args) {
+        int k = 0;
+        String outputFile = args[k++];
+        int numVertices = Integer.parseInt(args[k++]);
+        long numEdges = Long.parseLong(args[k++]);
+
+        double pA = Double.parseDouble(args[k++]);
+        double pB = Double.parseDouble(args[k++]);
+        double pC = Double.parseDouble(args[k++]);
+        double pD = Double.parseDouble(args[k++]);
+
+        System.out.println("Going to create graph with approx. " + numVertices + " vertices and " + numEdges + " edges");
+
+        GraphOutput outputInstance = null;
+        if (outputFile.startsWith("hdfs://")) outputInstance = new HDFSEdgeListOutput(outputFile);
+        else outputInstance = new EdgeListOutput(outputFile);
+
+        long t = System.currentTimeMillis();
+        RMATGraphGenerator generator = new RMATGraphGenerator(outputInstance,
+                pA, pB, pC, pD, numVertices, numEdges);
+        generator.execute();
+        System.out.println("Generating took " + (System.currentTimeMillis() - t) * 0.001 + " secs");
+
+    }
+
+}