Remove flag

cognitedata · Nov 2, 2023 · 9ee2fa5 · 9ee2fa5
1 parent 913920d
commit 9ee2fa5
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 34 deletions.
diff --git a/src/main/scala/cognite/spark/v1/AssetHierarchyBuilder.scala b/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

diff --git a/src/main/scala/cognite/spark/v1/DefaultSource.scala b/src/main/scala/cognite/spark/v1/DefaultSource.scala
@@ -406,9 +406,6 @@ object DefaultSource {
     val collectMetrics = toBoolean(parameters, "collectMetrics")
     val collectTestMetrics = toBoolean(parameters, "collectTestMetrics")
 
-    val enableSinglePartitionDeleteAssetHierarchy =
-      toBoolean(parameters, "enableSinglePartitionDeleteHierarchy", defaultValue = false)
-
     val saveMode = parseSaveMode(parameters)
     val parallelismPerPartition = {
       toPositiveInt(parameters, "parallelismPerPartition").getOrElse(
@@ -458,8 +455,7 @@ object DefaultSource {
       deleteMissingAssets = toBoolean(parameters, "deleteMissingAssets"),
       subtrees = subtreesOption,
       ignoreNullFields = toBoolean(parameters, "ignoreNullFields", defaultValue = true),
-      rawEnsureParent = toBoolean(parameters, "rawEnsureParent", defaultValue = true),
-      enableSinglePartitionDeleteAssetHierarchy = enableSinglePartitionDeleteAssetHierarchy
+      rawEnsureParent = toBoolean(parameters, "rawEnsureParent", defaultValue = true)
     )
   }
 

diff --git a/src/main/scala/cognite/spark/v1/RelationConfig.scala b/src/main/scala/cognite/spark/v1/RelationConfig.scala
@@ -21,8 +21,7 @@ final case class RelationConfig(
     deleteMissingAssets: Boolean,
     subtrees: AssetSubtreeOption,
     ignoreNullFields: Boolean,
-    rawEnsureParent: Boolean,
-    enableSinglePartitionDeleteAssetHierarchy: Boolean // flag to test whether single partition helps avoid NPE in asset hierarchy builder
+    rawEnsureParent: Boolean
 ) {
 
   /** Desired number of Spark partitions ~= partitions / parallelismPerPartition */

diff --git a/src/test/scala/cognite/spark/v1/SparkTest.scala b/src/test/scala/cognite/spark/v1/SparkTest.scala
@@ -251,8 +251,7 @@ trait SparkTest {
       deleteMissingAssets = false,
       subtrees = AssetSubtreeOption.Ingest,
       ignoreNullFields = true,
-      rawEnsureParent = false,
-      enableSinglePartitionDeleteAssetHierarchy = false
+      rawEnsureParent = false
     )
 
   private def getCounterSafe(metricName: String): Option[Long] =