added example on how to avoid map partitions

src/test/scala/com/johnsnowlabs/nlp/annotators/ner/dl/NerDLSpec.scala

@@ -24,6 +24,9 @@ import com.johnsnowlabs.nlp.util.io.ResourceHelper
 import com.johnsnowlabs.tags.{FastTest, SlowTest}
 import com.johnsnowlabs.util.{Benchmark, FileHelper}
 import org.apache.spark.ml.Pipeline
+import org.apache.spark.sql.{SparkSession, SparkSessionExtensions}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Expression, NamedExpression}
+import org.apache.spark.sql.connector.expressions.Expression
 import org.scalatest.flatspec.AnyFlatSpec
 
 import scala.io.Source
@@ -252,6 +255,73 @@ class NerDLSpec extends AnyFlatSpec {
 
   }
 
+  "NerDLModel" should "avoid mappartition" taggedAs FastTest in {
+
+    import org.apache.spark.sql.catalyst.plans.logical._
+    import org.apache.spark.sql.catalyst.rules.Rule
+
+    // Define a custom optimization rule
+    case class OptimizeMapPartitions(sparkSession: SparkSession) extends Rule[LogicalPlan] {
+      override def apply(plan: LogicalPlan): LogicalPlan = {
+
+        plan transformUp {
+          case mapPartitions@MapPartitions(func, objAttr, child) =>
+
+            // Does the plan look for the NER?
+            val callLambda = shouldCallLambda(mapPartitions, plan)
+            if (callLambda)
+              mapPartitions
+            else
+              child // we do nothing here!
+          case other => other // No change for other operations
+        }
+      }
+
+      private def shouldCallLambda(mapPartitions: MapPartitions, plan: LogicalPlan): Boolean = {
+        // things you have to do for money :)
+        var condition = false
+
+        plan transformUp {
+          // IMPORTANT:: here we determine whether the plan will require the NER or not!!
+          // child is the input to the projection, not used
+          case projection@Project(projectList, child) =>
+            // need to find a projection for which this mapPartition is present
+            // and the projection must include an NER column
+            if (projection.containsChild.head.children.contains(mapPartitions) && // must contain mapPartitions
+              projectList.exists(_.name.startsWith("ner")))
+              condition = true
+            projection
+        }
+        condition
+      }
+    }
+
+    // Register your Catalyst plugin
+    ResourceHelper.spark.experimental.extraOptimizations = Seq(OptimizeMapPartitions(ResourceHelper.spark))
+
+    val nerModel = NerDLModel
+      .pretrained()
+      .setInputCols("sentence", "token", "embeddings")
+      .setOutputCol("ner")
+
+    val conll = CoNLL(explodeSentences = false)
+    val training_data =
+      conll.readDataset(ResourceHelper.spark, "src/test/resources/conll2003/eng.train")
+
+    val embeddings = WordEmbeddingsModel.pretrained()
+
+    val pipeline = new Pipeline()
+      .setStages(Array(embeddings, nerModel))
+
+    //map Partitions not called
+    pipeline.fit(training_data).transform(training_data).select("token").collect()
+
+    //map Partitions called
+    pipeline.fit(training_data).transform(training_data).select("ner").collect()
+
+
+  }
+
   "NerDLApproach" should "benchmark test" taggedAs SlowTest in {
 
     val conll = CoNLL(explodeSentences = false)