Asset Hierarchy - Remove single partition deletion flag

src/main/scala/cognite/spark/v1/AssetHierarchyBuilder.scala

@@ -78,31 +78,26 @@ class AssetHierarchyBuilder(config: RelationConfig)(val sqlContext: SQLContext)
 
   import CdpConnector.ioRuntime
 
-  def delete(data: DataFrame): Unit = {
-    val partitionedData = if (config.enableSinglePartitionDeleteAssetHierarchy) {
-      data.repartition(numPartitions = 1)
-    } else {
-      data
-    }
-
-    partitionedData.foreachPartition((rows: Iterator[Row]) => {
-      val deletes = rows.map(r => fromRow[DeleteItemByCogniteId](r))
-      Stream
-        .fromIterator[IO](deletes, chunkSize = batchSize)
-        .chunks
-        .parEvalMapUnordered(config.parallelismPerPartition) { chunk =>
-          client.assets
-            .deleteRecursive(
-              chunk.toVector.map(_.toCogniteId),
-              recursive = true,
-              ignoreUnknownIds = true)
-            .flatTap(_ => incMetrics(itemsDeleted, chunk.size))
-        }
-        .compile
-        .drain
-        .unsafeRunSync()
-    })
-  }
+  def delete(data: DataFrame): Unit =
+    data
+      .repartition(numPartitions = 1)
+      .foreachPartition((rows: Iterator[Row]) => {
+        val deletes = rows.map(r => fromRow[DeleteItemByCogniteId](r))
+        Stream
+          .fromIterator[IO](deletes, chunkSize = batchSize)
+          .chunks
+          .parEvalMapUnordered(config.parallelismPerPartition) { chunk =>
+            client.assets
+              .deleteRecursive(
+                chunk.toVector.map(_.toCogniteId),
+                recursive = true,
+                ignoreUnknownIds = true)
+              .flatTap(_ => incMetrics(itemsDeleted, chunk.size))
+          }
+          .compile
+          .drain
+          .unsafeRunSync()
+      })
 
   def buildFromDf(data: DataFrame): Unit =
     // Do not use .collect to run the builder on one of the executors and not on the driver

src/main/scala/cognite/spark/v1/DefaultSource.scala

@@ -406,8 +406,8 @@ object DefaultSource {
     val collectMetrics = toBoolean(parameters, "collectMetrics")
     val collectTestMetrics = toBoolean(parameters, "collectTestMetrics")
 
-    val enableSinglePartitionDeleteAssetHierarchy =
-      toBoolean(parameters, "enableSinglePartitionDeleteHierarchy", defaultValue = false)
+    val userSinglePartition =
+      toBoolean(parameters, "useSinglePartition", defaultValue = false)
 
     val saveMode = parseSaveMode(parameters)
     val parallelismPerPartition = {
@@ -459,7 +459,7 @@ object DefaultSource {
       subtrees = subtreesOption,
       ignoreNullFields = toBoolean(parameters, "ignoreNullFields", defaultValue = true),
       rawEnsureParent = toBoolean(parameters, "rawEnsureParent", defaultValue = true),
-      enableSinglePartitionDeleteAssetHierarchy = enableSinglePartitionDeleteAssetHierarchy
+      useSinglePartition = userSinglePartition
     )
   }
 

src/main/scala/cognite/spark/v1/RelationConfig.scala

@@ -22,7 +22,7 @@ final case class RelationConfig(
     subtrees: AssetSubtreeOption,
     ignoreNullFields: Boolean,
     rawEnsureParent: Boolean,
-    enableSinglePartitionDeleteAssetHierarchy: Boolean // flag to test whether single partition helps avoid NPE in asset hierarchy builder
+    useSinglePartition: Boolean // a flag for helping to test NPEs by moving the data frame to a single partition
 ) {
 
   /** Desired number of Spark partitions ~= partitions / parallelismPerPartition */

src/test/scala/cognite/spark/v1/SparkTest.scala

@@ -252,7 +252,7 @@ trait SparkTest {
       subtrees = AssetSubtreeOption.Ingest,
       ignoreNullFields = true,
       rawEnsureParent = false,
-      enableSinglePartitionDeleteAssetHierarchy = false
+      useSinglePartition = false
     )
 
   private def getCounterSafe(metricName: String): Option[Long] =