From 339fdb6481b0e45291338ac154ce6a272a4e71f7 Mon Sep 17 00:00:00 2001
From: Ben Chand <benchand1995@gmail.com>
Date: Thu, 2 Nov 2023 11:02:32 +0000
Subject: [PATCH] Rename flag

---
 src/main/scala/cognite/spark/v1/DefaultSource.scala  | 6 +++++-
 src/main/scala/cognite/spark/v1/RelationConfig.scala | 3 ++-
 src/test/scala/cognite/spark/v1/SparkTest.scala      | 3 ++-
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/cognite/spark/v1/DefaultSource.scala b/src/main/scala/cognite/spark/v1/DefaultSource.scala
index ccfdc4449..666a5b40d 100644
--- 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
index cbaeb3db9..29cb833a2 100644
--- 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
index f1455653a..153c89b93 100644
--- 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] =