Merge pull request #323 from ldbc/file-uri-fixes

dszakallas · web-flow · commit 36d6628e7717 · 2021-08-10T19:17:49.000+02:00
File uri fixes
diff --git a/pom.xml b/pom.xml
@@ -72,11 +72,6 @@
       <artifactId>shapeless_${scala.compat.version}</artifactId>
       <version>2.3.3</version>
     </dependency>
-    <dependency>
-      <groupId>com.github.pathikrit</groupId>
-      <artifactId>better-files_${scala.compat.version}</artifactId>
-      <version>3.8.0</version>
-    </dependency>
     <dependency>
       <groupId>com.github.scopt</groupId>
       <artifactId>scopt_${scala.compat.version}</artifactId>
diff --git a/src/main/java/ldbc/snb/datagen/hadoop/serializer/HadoopStaticSerializer.java b/src/main/java/ldbc/snb/datagen/hadoop/serializer/HadoopStaticSerializer.java
@@ -47,6 +47,7 @@ Linked Data Benchmark Council (http://www.ldbcouncil.org)
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
 
+import java.net.URI;
 import java.util.Iterator;
 import java.util.Set;
 import java.util.TreeSet;
@@ -71,7 +72,7 @@ public HadoopStaticSerializer(GeneratorConfiguration conf, Configuration hadoopC
 
     public void run() {
         try {
-            FileSystem fs = FileSystem.get(hadoopConf);
+            FileSystem fs = FileSystem.get(new URI(conf.getOutputDir()), hadoopConf);
             staticSerializer = new StaticSerializer[numPartitions];
             for (int i = 0; i < numPartitions; ++i) {
                 staticSerializer[i] = new StaticSerializer();
diff --git a/src/main/java/ldbc/snb/datagen/hadoop/sorting/HadoopCreationTimeSorter.java b/src/main/java/ldbc/snb/datagen/hadoop/sorting/HadoopCreationTimeSorter.java
diff --git a/src/main/java/ldbc/snb/datagen/hadoop/sorting/HadoopDeletionTimeSorter.java b/src/main/java/ldbc/snb/datagen/hadoop/sorting/HadoopDeletionTimeSorter.java
diff --git a/src/main/java/ldbc/snb/datagen/hadoop/sorting/SortPartitioner.java b/src/main/java/ldbc/snb/datagen/hadoop/sorting/SortPartitioner.java
diff --git a/src/main/java/ldbc/snb/datagen/hadoop/sorting/SortReducer.java b/src/main/java/ldbc/snb/datagen/hadoop/sorting/SortReducer.java
diff --git a/src/main/scala/ldbc/snb/datagen/factors/io/package.scala b/src/main/scala/ldbc/snb/datagen/factors/io/package.scala
@@ -3,7 +3,7 @@ package ldbc.snb.datagen.factors
 import ldbc.snb.datagen.io.{PathComponent, Writer}
 import ldbc.snb.datagen.util.Logging
 import ldbc.snb.datagen.model.{GraphLike, Mode}
-import better.files._
+import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.io.dataframes.DataFrameSink
 
 package object io {
@@ -16,7 +16,7 @@ package object io {
     override type CoRet = FactorTable[M]
 
     override def write(self: FactorTable[M], sink: FactorTableSink): Unit = {
-      val p = (sink.path / "factors" / sink.format / PathComponent[GraphLike[M]].path(self.source) / self.name).toString()
+      val p = (sink.path / "factors" / sink.format / PathComponent[GraphLike[M]].path(self.source) / self.name).toString
       self.data.coalesce(1).write(DataFrameSink(p, sink.format))
       log.info(s"Factor table ${self.name} written")
     }
diff --git a/src/main/scala/ldbc/snb/datagen/generation/serializer/SparkActivitySerializer.scala b/src/main/scala/ldbc/snb/datagen/generation/serializer/SparkActivitySerializer.scala
@@ -12,6 +12,7 @@ import org.apache.spark.TaskContext
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.SparkSession
 
+import java.net.URI
 import java.util
 import java.util.Collections
 import java.util.function.Consumer
@@ -35,7 +36,7 @@ object SparkActivitySerializer {
       val hadoopConf = serializableHadoopConf.value
       val buildDir = conf.getOutputDir
 
-      val fs = FileSystem.get(hadoopConf)
+      val fs = FileSystem.get(new URI(buildDir), hadoopConf)
       fs.mkdirs(new Path(buildDir))
 
       val dynamicActivitySerializer = new DynamicActivitySerializer()
diff --git a/src/main/scala/ldbc/snb/datagen/generation/serializer/SparkPersonSerializer.scala b/src/main/scala/ldbc/snb/datagen/generation/serializer/SparkPersonSerializer.scala
@@ -10,6 +10,8 @@ import org.apache.spark.TaskContext
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.SparkSession
 
+import java.net.URI
+
 object SparkPersonSerializer {
 
   def apply(
@@ -28,7 +30,7 @@ object SparkPersonSerializer {
         val partitionId = TaskContext.getPartitionId()
         val buildDir = conf.getOutputDir
 
-        val fs = FileSystem.get(hadoopConf)
+        val fs = FileSystem.get(new URI(buildDir), hadoopConf)
         fs.mkdirs(new Path(buildDir))
 
         dynamicPersonSerializer.initialize(
diff --git a/src/main/scala/ldbc/snb/datagen/io/dataframes.scala b/src/main/scala/ldbc/snb/datagen/io/dataframes.scala
@@ -1,6 +1,6 @@
 package ldbc.snb.datagen.io
 
-import ldbc.snb.datagen.syntax.fluentSyntaxOps
+import ldbc.snb.datagen.syntax._
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}
 
diff --git a/src/main/scala/ldbc/snb/datagen/io/graphs.scala b/src/main/scala/ldbc/snb/datagen/io/graphs.scala
@@ -1,10 +1,10 @@
 package ldbc.snb.datagen.io
 
-import better.files._
 import ldbc.snb.datagen.io.dataframes.{DataFrameSink, DataFrameSource}
 import ldbc.snb.datagen.model.EntityType.{Attr, Edge, Node}
 import ldbc.snb.datagen.model.Mode.Raw
 import ldbc.snb.datagen.model._
+import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.util.{Logging, SparkUI}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}
@@ -90,7 +90,7 @@ object graphs {
       TreeMap(self.entities.toSeq: _*).foreach {
         case (tpe, dataset) =>
           SparkUI.job(getClass.getSimpleName, s"write $tpe") {
-            val p = (sink.path / sink.format / PathComponent[GraphLike[M]].path(self) / tpe.entityPath).toString()
+            val p = (sink.path / sink.format / PathComponent[GraphLike[M]].path(self) / tpe.entityPath).toString
             log.info(s"$tpe: Writing started")
             val opts = getFormatOptions(sink.format, self.mode)
             the(dataset).write(DataFrameSink(p, sink.format, opts, SaveMode.Ignore))
@@ -113,7 +113,7 @@ object graphs {
       TreeMap(self.entities.mapValues(ev).toSeq: _*).foreach {
         case (tpe, BatchedEntity(snapshot, insertBatches, deleteBatches)) =>
           SparkUI.job(getClass.getSimpleName, s"write $tpe snapshot") {
-            val p = (sink.path / sink.format / PathComponent[GraphLike[M]].path(self) / "initial_snapshot" / tpe.entityPath).toString()
+            val p = (sink.path / sink.format / PathComponent[GraphLike[M]].path(self) / "initial_snapshot" / tpe.entityPath).toString
             log.info(s"$tpe: Writing snapshot")
             snapshot.write(DataFrameSink(p, sink.format, opts, SaveMode.Ignore))
             log.info(s"$tpe: Writing snapshot completed")
@@ -123,7 +123,7 @@ object graphs {
             batches.foreach {
               case Batched(entity, partitionKeys) =>
                 SparkUI.job(getClass.getSimpleName, s"write $tpe $operation") {
-                  val p = (sink.path / sink.format / PathComponent[GraphLike[M]].path(self) / operation / tpe.entityPath).toString()
+                  val p = (sink.path / sink.format / PathComponent[GraphLike[M]].path(self) / operation / tpe.entityPath).toString
                   log.info(f"$tpe: Writing $operation")
                   entity.write(DataFrameSink(p, sink.format, opts, SaveMode.Ignore, partitionBy = partitionKeys))
                   log.info(f"$tpe: Writing $operation completed")
diff --git a/src/main/scala/ldbc/snb/datagen/spark/LdbcDatagen.scala b/src/main/scala/ldbc/snb/datagen/spark/LdbcDatagen.scala
@@ -50,7 +50,8 @@ object LdbcDatagen extends SparkApp {
 
       opt[String]('o', "output-dir")
         .action((x, c) => args.outputDir.set(c)(x))
-        .text("output directory")
+        .text("path on the cluster filesystem, where Datagen outputs. Can be a URI (e.g S3, ADLS, HDFS) or a " +
+          "path in which case the default cluster file system is used.")
 
       opt[Int]('n', "num-threads")
         .action((x, c) => args.numThreads.set(c)(Some(x)))
diff --git a/src/main/scala/ldbc/snb/datagen/syntax/PathSyntax.scala b/src/main/scala/ldbc/snb/datagen/syntax/PathSyntax.scala
@@ -0,0 +1,46 @@
+package ldbc.snb.datagen.syntax
+
+import org.apache.hadoop.fs.Path
+
+import java.net.URI
+import scala.language.implicitConversions
+
+trait PathSyntax {
+  @`inline` implicit final def pathSyntaxOpsForString[A](a: String): PathSyntaxOpsForString = new PathSyntaxOpsForString(a)
+  @`inline` implicit final def pathSyntaxOpsForPath[A](a: Path): PathSyntaxOpsForPath = new PathSyntaxOpsForPath(a)
+  @`inline` implicit final def pathSyntaxOpsForUri[A](a: URI): PathSyntaxOpsForUri = new PathSyntaxOpsForUri(a)
+}
+
+final class PathSyntaxOpsForString(private val self: String) extends AnyVal {
+  import PathSyntaxOpsHelpers._
+  def /(child: String): Path = join(new Path(self), new Path(child))
+  def /(child: Path): Path = join(new Path(self), child)
+}
+
+final class PathSyntaxOpsForPath(private val self: Path) extends AnyVal {
+  import PathSyntaxOpsHelpers._
+  def /(child: String): Path = join(self, new Path(child))
+  def /(child: Path): Path = join(self, child)
+}
+
+final class PathSyntaxOpsForUri(private val self: URI) extends AnyVal {
+  import PathSyntaxOpsHelpers._
+  def /(child: String): Path = join(new Path(self), new Path(child))
+  def /(child: Path): Path = join(new Path(self), child)
+}
+
+private[syntax] object PathSyntaxOpsHelpers {
+  def join(path1: Path, path2: Path): Path = new Path(ensureTrailingSlashForAbsoluteUri(path1), path2)
+
+  private[this] def ensureTrailingSlashForAbsoluteUri(path: Path): Path = {
+    if (path.isAbsolute)
+      return path
+
+    val uri = path.toUri
+
+    if (uri.getScheme == null || uri.getPath != "")
+      return path
+
+    new Path(new URI(uri.getScheme, uri.getAuthority, "/", uri.getQuery, uri.getFragment))
+  }
+}
diff --git a/src/main/scala/ldbc/snb/datagen/syntax/package.scala b/src/main/scala/ldbc/snb/datagen/syntax/package.scala
@@ -5,3 +5,4 @@ package object syntax
     with UseSyntax
     with SparkSqlSyntax
     with BooleanSyntax
+    with PathSyntax
