Rename flag

cognitedata · Nov 2, 2023 · 339fdb6 · 339fdb6
1 parent 9ee2fa5
commit 339fdb6
Show file tree

Hide file tree

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

diff --git a/src/main/scala/cognite/spark/v1/RelationConfig.scala b/src/main/scala/cognite/spark/v1/RelationConfig.scala
@@ -21,7 +21,8 @@ final case class RelationConfig(
     deleteMissingAssets: Boolean,
     subtrees: AssetSubtreeOption,
     ignoreNullFields: Boolean,
-    rawEnsureParent: Boolean
+    rawEnsureParent: Boolean,
+    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 */

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