Just 1 partition code for now

cognitedata · Oct 30, 2023 · ecb2e8b · ecb2e8b
1 parent d22ee7b
commit ecb2e8b
Showing 1 changed file with 23 additions and 33 deletions.
diff --git a/src/main/scala/cognite/spark/v1/AssetHierarchyBuilder.scala b/src/main/scala/cognite/spark/v1/AssetHierarchyBuilder.scala
@@ -78,42 +78,32 @@ class AssetHierarchyBuilder(config: RelationConfig)(val sqlContext: SQLContext)
 
   import CdpConnector.ioRuntime
 
-  def delete(data: DataFrame): Unit =
-    if (config.enableSinglePartitionDeleteAssetHierarchy) {
-      data
-        .repartition(numPartitions = 1)
-        .foreachPartition((rows: Iterator[Row]) => {
-          batchedOperation[DeleteItemByCogniteId, Unit](
-            rows.map(r => fromRow[DeleteItemByCogniteId](r)).toVector,
-            // The API calls throw exception when any of the ids do not exist
-            ids =>
-              client.assets
-                .deleteRecursive(ids.map(_.toCogniteId), recursive = true, ignoreUnknownIds = true)
-                .flatTap(_ => incMetrics(itemsDeleted, ids.size))
-                .map(Seq(_)),
-            batchSize = batchSize
-          ).void.unsafeRunSync()
-        })
+  def delete(data: DataFrame): Unit = {
+    val partitionedData = if (config.enableSinglePartitionDeleteAssetHierarchy) {
+      data.repartition(numPartitions = 1)
     } else {
-      data.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()
-      })
+      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 buildFromDf(data: DataFrame): Unit =
     // Do not use .collect to run the builder on one of the executors and not on the driver
     data