added write options to set numberOfShards and collectionType

Author: rashtao

README.md

@@ -203,6 +203,8 @@ according to the related target collection definition and is different from the
 ### Write Configuration
 
 - `table`: target ArangoDB collection name (required)
+- `table.shards`: number of shards of the created collection (in case of SaveMode `Append` or `Overwrite`)
+- `table.type`: type (`document`|`edge`) of the created collection (in case of SaveMode `Append` or `Overwrite`)
 - `batch.size`: writing batch size, default `1000`
 - `wait.sync`: whether to wait until the documents have been synced to disk (`true`|`false`)
 - `confirm.truncate`: confirm to truncate table when using `SaveMode.Overwrite` mode, default `false`

arangodb-spark-commons/src/main/scala/org/apache/spark/sql/arangodb/commons/ArangoClient.scala

@@ -99,16 +99,16 @@ class ArangoClient(options: ArangoOptions) {
     .collection(options.writeOptions.collection)
     .exists()
 
-  def createCollection(): Unit = arangoDB
-    .db(options.writeOptions.db)
-    .collection(options.writeOptions.collection)
-    .create(new CollectionCreateOptions()
-      // TODO:
-      //      .`type`()
-      //      .numberOfShards()
-      //      .replicationFactor()
-      //      .minReplicationFactor()
-    )
+  def createCollection(): Unit = {
+    val opts = new CollectionCreateOptions()
+    options.writeOptions.numberOfShards.foreach(opts.numberOfShards(_))
+    options.writeOptions.collectionType.foreach(ct => opts.`type`(ct.get()))
+
+    arangoDB
+      .db(options.writeOptions.db)
+      .collection(options.writeOptions.collection)
+      .create(opts)
+  }
 
   def truncate(): Unit = arangoDB
     .db(options.writeOptions.db)

arangodb-spark-commons/src/main/scala/org/apache/spark/sql/arangodb/commons/ArangoOptions.scala

@@ -20,7 +20,7 @@
 
 package org.apache.spark.sql.arangodb.commons
 
-import com.arangodb.ArangoDB
+import com.arangodb.{ArangoDB, entity}
 import com.arangodb.model.OverwriteMode
 
 import java.io.ByteArrayInputStream
@@ -87,6 +87,8 @@ object ArangoOptions {
   val STREAM = "stream"
 
   // write options
+  val NUMBER_OF_SHARDS = "table.shards"
+  val COLLECTION_TYPE = "table.type"
   val WAIT_FOR_SYNC = "wait.sync"
   val CONFIRM_TRUNCATE = "confirm.truncate"
   val OVERWRITE_MODE = "overwrite.mode"
@@ -178,6 +180,8 @@ class ArangoReadOptions(options: Map[String, String]) extends CommonOptions(opti
 class ArangoWriteOptions(options: Map[String, String]) extends CommonOptions(options) {
   val batchSize: Int = options.get(ArangoOptions.BATCH_SIZE).map(_.toInt).getOrElse(1000)
   val collection: String = getRequired(ArangoOptions.COLLECTION)
+  val numberOfShards: Option[Int] = options.get(ArangoOptions.NUMBER_OF_SHARDS).map(_.toInt)
+  val collectionType: Option[CollectionType] = options.get(ArangoOptions.COLLECTION_TYPE).map(CollectionType(_))
   val waitForSync: Option[Boolean] = options.get(ArangoOptions.WAIT_FOR_SYNC).map(_.toBoolean)
   val confirmTruncate: Boolean = options.getOrElse(ArangoOptions.CONFIRM_TRUNCATE, "false").toBoolean
   val overwriteMode: Option[OverwriteMode] = options.get(ArangoOptions.OVERWRITE_MODE).map(OverwriteMode.valueOf)
@@ -226,3 +230,23 @@ object Protocol {
     case _ => throw new IllegalArgumentException(s"${ArangoOptions.PROTOCOL}: $value")
   }
 }
+
+sealed trait CollectionType {
+  def get(): entity.CollectionType
+}
+
+object CollectionType {
+  case object DOCUMENT extends CollectionType {
+    override def get() = entity.CollectionType.DOCUMENT
+  }
+
+  case object EDGE extends CollectionType {
+    override def get() = entity.CollectionType.EDGES
+  }
+
+  def apply(value: String): CollectionType = value match {
+    case "document" => DOCUMENT
+    case "edge" => EDGE
+    case _ => throw new IllegalArgumentException(s"${ArangoOptions.COLLECTION_TYPE}: $value")
+  }
+}

integration-tests/src/test/scala/org/apache/spark/sql/arangodb/datasource/write/CreateCollectionTest.scala

@@ -0,0 +1,77 @@
+package org.apache.spark.sql.arangodb.datasource.write
+
+import com.arangodb.ArangoCollection
+import com.arangodb.entity.CollectionType
+import org.apache.spark.sql.SaveMode
+import org.apache.spark.sql.arangodb.commons.ArangoOptions
+import org.apache.spark.sql.arangodb.datasource.BaseSparkTest
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.MethodSource
+
+
+class CreateCollectionTest extends BaseSparkTest {
+
+  private val collectionName = "chessPlayersCreateCollection"
+  private val collection: ArangoCollection = db.collection(collectionName)
+
+  import spark.implicits._
+
+  private val df = Seq(
+    ("a/1", "b/1"),
+    ("a/2", "b/2"),
+    ("a/3", "b/3"),
+    ("a/4", "b/4"),
+    ("a/5", "b/5"),
+    ("a/6", "b/6")
+  ).toDF("_from", "_to")
+    .repartition(3)
+
+  @BeforeEach
+  def beforeEach(): Unit = {
+    if (collection.exists()) {
+      collection.drop()
+    }
+  }
+
+  @ParameterizedTest
+  @MethodSource(Array("provideProtocolAndContentType"))
+  def saveModeAppend(protocol: String, contentType: String): Unit = {
+    df.write
+      .format(BaseSparkTest.arangoDatasource)
+      .mode(SaveMode.Append)
+      .options(options + (
+        ArangoOptions.COLLECTION -> collectionName,
+        ArangoOptions.PROTOCOL -> protocol,
+        ArangoOptions.CONTENT_TYPE -> contentType,
+        ArangoOptions.NUMBER_OF_SHARDS -> "5",
+        ArangoOptions.COLLECTION_TYPE -> "edge"
+      ))
+      .save()
+
+    assertThat(collection.getProperties.getNumberOfShards).isEqualTo(5)
+    assertThat(collection.getProperties.getType.getType).isEqualTo(CollectionType.EDGES.getType)
+  }
+
+  @ParameterizedTest
+  @MethodSource(Array("provideProtocolAndContentType"))
+  def saveModeOverwrite(protocol: String, contentType: String): Unit = {
+    df.write
+      .format(BaseSparkTest.arangoDatasource)
+      .mode(SaveMode.Overwrite)
+      .options(options + (
+        ArangoOptions.COLLECTION -> collectionName,
+        ArangoOptions.PROTOCOL -> protocol,
+        ArangoOptions.CONTENT_TYPE -> contentType,
+        ArangoOptions.CONFIRM_TRUNCATE -> "true",
+        ArangoOptions.NUMBER_OF_SHARDS -> "5",
+        ArangoOptions.COLLECTION_TYPE -> "edge"
+      ))
+      .save()
+
+    assertThat(collection.getProperties.getNumberOfShards).isEqualTo(5)
+    assertThat(collection.getProperties.getType.getType).isEqualTo(CollectionType.EDGES.getType)
+  }
+
+}