[SPARKNLP-1163] Adding title chunking strategy (#14594)

danilojsl · DevinTDHa · commit 3e0aa61482e1 · 2025-06-09T11:10:05.000+02:00
diff --git a/python/sparknlp/partition/partition_properties.py b/python/sparknlp/partition/partition_properties.py
@@ -256,7 +256,7 @@ def setThreshold(self, value):
     def getThreshold(self):
         return self.getOrDefault(self.threshold)
 
-class HasSemanticChunkerProperties(Params):
+class HasChunkerProperties(Params):
 
     chunkingStrategy = Param(
         Params._dummy(),
@@ -296,4 +296,24 @@ def setNewAfterNChars(self, value):
     )
 
     def setOverlap(self, value):
-        return self._set(overlap=value)
+        return self._set(overlap=value)
+
+    combineTextUnderNChars = Param(
+        Params._dummy(),
+        "combineTextUnderNChars",
+        "Threshold to merge adjacent small sections",
+        typeConverter=TypeConverters.toInt
+    )
+
+    def setCombineTextUnderNChars(self, value):
+        return self._set(combineTextUnderNChars=value)
+
+    overlapAll = Param(
+        Params._dummy(),
+        "overlapAll",
+        "Apply overlap context between all sections, not just split chunks",
+        typeConverter=TypeConverters.toBoolean
+    )
+
+    def setOverlapAll(self, value):
+        return self._set(overlapAll=value)
diff --git a/python/sparknlp/partition/partition_transformer.py b/python/sparknlp/partition/partition_transformer.py
@@ -23,7 +23,7 @@ class PartitionTransformer(
     HasHTMLReaderProperties,
     HasPowerPointProperties,
     HasTextReaderProperties,
-    HasSemanticChunkerProperties
+    HasChunkerProperties
 ):
     """
     The PartitionTransformer annotator allows you to use the Partition feature more smoothly
@@ -194,5 +194,7 @@ def __init__(self, classname="com.johnsnowlabs.partition.PartitionTransformer",
             chunkingStrategy="",
             maxCharacters=100,
             newAfterNChars=-1,
-            overlap=0
+            overlap=0,
+            combineTextUnderNChars=0,
+            overlapAll=False
         )
diff --git a/src/main/scala/com/johnsnowlabs/partition/BasicChunker.scala b/src/main/scala/com/johnsnowlabs/partition/BasicChunker.scala
@@ -19,10 +19,6 @@ import com.johnsnowlabs.reader.HTMLElement
 
 import scala.collection.mutable
 
-case class Chunk(elements: List[HTMLElement]) {
-  def length: Int = elements.map(_.content.length).sum
-}
-
 object BasicChunker {
 
   /** Splits a list of [[HTMLElement]]s into chunks constrained by a maximum number of characters.
diff --git a/src/main/scala/com/johnsnowlabs/partition/Chunk.scala b/src/main/scala/com/johnsnowlabs/partition/Chunk.scala
@@ -0,0 +1,7 @@
+package com.johnsnowlabs.partition
+
+import com.johnsnowlabs.reader.HTMLElement
+
+case class Chunk(elements: List[HTMLElement]) {
+  def length: Int = elements.map(_.content.length).sum
+}
diff --git a/src/main/scala/com/johnsnowlabs/partition/HasChunkerProperties.scala b/src/main/scala/com/johnsnowlabs/partition/HasChunkerProperties.scala
@@ -18,7 +18,7 @@ package com.johnsnowlabs.partition
 import com.johnsnowlabs.nlp.ParamsAndFeaturesWritable
 import org.apache.spark.ml.param.Param
 
-trait HasSemanticChunkerProperties extends ParamsAndFeaturesWritable {
+trait HasChunkerProperties extends ParamsAndFeaturesWritable {
 
   val chunkingStrategy = new Param[String](this, "chunkingStrategy", "Set the chunking strategy")
 
@@ -39,6 +39,26 @@ trait HasSemanticChunkerProperties extends ParamsAndFeaturesWritable {
 
   def setOverlap(value: Int): this.type = set(overlap, value)
 
-  setDefault(chunkingStrategy -> "", maxCharacters -> 100, newAfterNChars -> -1, overlap -> 0)
+  val combineTextUnderNChars =
+    new Param[Int](this, "combineTextUnderNChars", "Threshold to merge adjacent small sections")
+
+  def setComBineTextUnderNChars(value: Int): this.type =
+    set(combineTextUnderNChars, value)
+
+  val overlapAll =
+    new Param[Boolean](
+      this,
+      "overlapAll",
+      "Apply overlap context between all sections, not just split chunks")
+
+  def setOverlapAll(value: Boolean): this.type = set(overlapAll, value)
+
+  setDefault(
+    chunkingStrategy -> "",
+    maxCharacters -> 100,
+    newAfterNChars -> -1,
+    overlap -> 0,
+    combineTextUnderNChars -> 0,
+    overlapAll -> false)
 
 }
diff --git a/src/main/scala/com/johnsnowlabs/partition/Partition.scala b/src/main/scala/com/johnsnowlabs/partition/Partition.scala
@@ -146,7 +146,7 @@ class Partition(params: java.util.Map[String, String] = new java.util.HashMap())
 
     val partitionResult = reader(path)
     if (hasChunkerStrategy) {
-      val chunker = new SemanticChunker(params.asScala.toMap)
+      val chunker = new PartitionChunker(params.asScala.toMap)
       partitionResult.withColumn(
         "chunks",
         chunker.chunkUDF()(partitionResult(sparkNLPReader.getOutputColumn)))
diff --git a/src/main/scala/com/johnsnowlabs/partition/PartitionChunker.scala b/src/main/scala/com/johnsnowlabs/partition/PartitionChunker.scala
@@ -16,15 +16,20 @@
 package com.johnsnowlabs.partition
 
 import com.johnsnowlabs.partition.BasicChunker.chunkBasic
+import com.johnsnowlabs.partition.TitleChunker.chunkByTitle
 import com.johnsnowlabs.reader.HTMLElement
-import com.johnsnowlabs.reader.util.PartitionOptions.{getDefaultInt, getDefaultString}
+import com.johnsnowlabs.reader.util.PartitionOptions.{
+  getDefaultBoolean,
+  getDefaultInt,
+  getDefaultString
+}
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.expressions.UserDefinedFunction
 import org.apache.spark.sql.functions.udf
 
 import scala.collection.mutable
 
-class SemanticChunker(chunkerOptions: Map[String, String]) extends Serializable {
+class PartitionChunker(chunkerOptions: Map[String, String]) extends Serializable {
 
   def chunkUDF(): UserDefinedFunction = {
     udf((elements: Seq[Row]) => {
@@ -37,6 +42,14 @@ class SemanticChunker(chunkerOptions: Map[String, String]) extends Serializable
 
       val chunks = getChunkerStrategy match {
         case "basic" => chunkBasic(htmlElements, getMaxCharacters, getNewAfterNChars, getOverlap)
+        case "byTitle" | "by_title" =>
+          chunkByTitle(
+            htmlElements,
+            getMaxCharacters,
+            getCombineTextUnderNChars,
+            getOverlap,
+            getNewAfterNChars,
+            getOverlapAll)
         case _ =>
           throw new IllegalArgumentException(s"Unknown chunker strategy: $getChunkerStrategy")
       }
@@ -64,4 +77,15 @@ class SemanticChunker(chunkerOptions: Map[String, String]) extends Serializable
       default = "none")
   }
 
+  private def getCombineTextUnderNChars: Int = {
+    getDefaultInt(
+      chunkerOptions,
+      Seq("combineTextUnderNChars", "combine_text_under_n_chars"),
+      default = 0)
+  }
+
+  private def getOverlapAll: Boolean = {
+    getDefaultBoolean(chunkerOptions, Seq("overlapAll", "overlap_all"), default = false)
+  }
+
 }
diff --git a/src/main/scala/com/johnsnowlabs/partition/PartitionTransformer.scala b/src/main/scala/com/johnsnowlabs/partition/PartitionTransformer.scala
@@ -86,7 +86,7 @@ class PartitionTransformer(override val uid: String)
     with HasPowerPointProperties
     with HasTextReaderProperties
     with HasPdfProperties
-    with HasSemanticChunkerProperties {
+    with HasChunkerProperties {
 
   def this() = this(Identifiable.randomUID("PartitionTransformer"))
   protected val logger: Logger = LoggerFactory.getLogger(getClass.getName)
@@ -155,7 +155,9 @@ class PartitionTransformer(override val uid: String)
       "chunkingStrategy" -> $(chunkingStrategy),
       "maxCharacters" -> $(maxCharacters).toString,
       "newAfterNChars" -> $(newAfterNChars).toString,
-      "overlap" -> $(overlap).toString)
+      "overlap" -> $(overlap).toString,
+      "combineTextUnderNChars" -> $(combineTextUnderNChars).toString,
+      "overlapAll" -> $(overlapAll).toString)
     val partitionInstance = new Partition(params.asJava)
 
     val inputColum = if (get(inputCols).isDefined) {
diff --git a/src/main/scala/com/johnsnowlabs/partition/TitleChunker.scala b/src/main/scala/com/johnsnowlabs/partition/TitleChunker.scala
@@ -0,0 +1,151 @@
+/*
+ * Copyright 2017-2025 John Snow Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.johnsnowlabs.partition
+
+import com.johnsnowlabs.reader.{ElementType, HTMLElement}
+
+import scala.collection.mutable
+
+object TitleChunker {
+
+  /** Splits a list of HTML elements into semantically grouped Chunks based on Title and Table
+    * markers.
+    *
+    * @param elements
+    *   List of input HTML elements to chunk.
+    * @param maxCharacters
+    *   Maximum length allowed per chunk. Longer sections are split.
+    * @param combineTextUnderNChars
+    *   Threshold to merge adjacent small sections.
+    * @param overlap
+    *   Number of characters to repeat between consecutive chunks.
+    * @param newAfterNChars
+    *   Soft limit to trigger new section if length exceeded, even before maxCharacters.
+    * @param overlapAll
+    *   Apply overlap context between all sections, not just split chunks.
+    * @return
+    *   List of Chunks partitioned by title and content heuristics.
+    */
+  def chunkByTitle(
+      elements: List[HTMLElement],
+      maxCharacters: Int,
+      combineTextUnderNChars: Int = 0,
+      overlap: Int = 0,
+      newAfterNChars: Int = -1,
+      overlapAll: Boolean = false): List[Chunk] = {
+
+    val softLimit = if (newAfterNChars <= 0) maxCharacters else newAfterNChars
+    val chunks = mutable.ListBuffer.empty[Chunk]
+    val sections = mutable.ListBuffer.empty[List[HTMLElement]]
+    var currentSection = List.empty[HTMLElement]
+    var currentLength = 0
+    var currentPage = -1
+
+    for (element <- elements) {
+      val elementLength = element.content.length
+      val isTable = element.elementType == "Table"
+      val elementPage = element.metadata.getOrElse("pageNumber", "-1").toInt
+
+      val pageChanged = currentPage != -1 && elementPage != currentPage
+      val softLimitExceeded = currentSection.length >= 2 &&
+        (currentLength + elementLength > softLimit)
+
+      if (isTable) {
+        if (currentSection.nonEmpty) sections += currentSection
+        sections += List(element)
+        currentSection = List.empty
+        currentLength = 0
+        currentPage = -1
+      } else if (pageChanged || softLimitExceeded) {
+        if (currentSection.nonEmpty) sections += currentSection
+        currentSection = List(element)
+        currentLength = elementLength
+        currentPage = elementPage
+      } else {
+        currentSection :+= element
+        currentLength += elementLength
+        currentPage = elementPage
+      }
+    }
+    if (currentSection.nonEmpty) sections += currentSection
+
+    val mergedSections = sections.foldLeft(List.empty[List[HTMLElement]]) { (acc, section) =>
+      val sectionLength = section.map(_.content.length).sum
+      val canMerge = combineTextUnderNChars > 0 &&
+        sectionLength < combineTextUnderNChars &&
+        acc.nonEmpty &&
+        acc.last.exists(_.elementType != "Table") &&
+        section.exists(_.elementType != "Table")
+
+      if (canMerge) {
+        acc.init :+ (acc.last ++ section)
+      } else {
+        acc :+ section
+      }
+    }
+
+    var lastNarrativeText = ""
+    for (section <- mergedSections) {
+      if (section.exists(_.elementType == "Table")) {
+        chunks += Chunk(section)
+        lastNarrativeText = ""
+      } else {
+        val sectionText = section.map(_.content).mkString(" ")
+        val content =
+          if (overlap > 0 && lastNarrativeText.nonEmpty && (overlapAll || sectionText.length > maxCharacters))
+            lastNarrativeText.takeRight(overlap) + " " + sectionText
+          else sectionText
+
+        val merged = HTMLElement(ElementType.NARRATIVE_TEXT, content.trim, section.head.metadata)
+        val split = if (content.length > maxCharacters) {
+          splitHTMLElement(merged, maxCharacters, overlap)
+        } else List(merged)
+
+        chunks ++= split.map(e => Chunk(List(e)))
+        lastNarrativeText = sectionText
+      }
+    }
+
+    chunks.toList
+  }
+
+  private def splitHTMLElement(
+      element: HTMLElement,
+      maxLen: Int,
+      overlap: Int): List[HTMLElement] = {
+
+    val words = element.content.split(" ")
+    val buffer = mutable.ListBuffer.empty[HTMLElement]
+    var chunk = new StringBuilder
+
+    for (word <- words) {
+      if (chunk.length + word.length + 1 > maxLen) {
+        val text = chunk.toString().trim
+        buffer += element.copy(content = text)
+        chunk = new StringBuilder
+        if (overlap > 0 && text.length >= overlap)
+          chunk.append(text.takeRight(overlap)).append(" ")
+      }
+      chunk.append(word).append(" ")
+    }
+
+    if (chunk.nonEmpty)
+      buffer += element.copy(content = chunk.toString().trim)
+
+    buffer.toList
+  }
+
+}
diff --git a/src/test/scala/com/johnsnowlabs/partition/PartitionChunkerTest.scala b/src/test/scala/com/johnsnowlabs/partition/PartitionChunkerTest.scala
@@ -23,6 +23,7 @@ import org.scalatest.flatspec.AnyFlatSpec
 class PartitionChunkerTest extends AnyFlatSpec {
   import ResourceHelper.spark.implicits._
   val txtDirectory = "src/test/resources/reader/txt"
+  val htmlDirectory = "src/test/resources/reader/html"
 
   "Partition" should "perform basic chunk text" taggedAs FastTest in {
     val partitionOptions = Map("contentType" -> "text/plain", "chunkingStrategy" -> "basic")
@@ -38,4 +39,17 @@ class PartitionChunkerTest extends AnyFlatSpec {
     assert(chunkDf.count() > 1)
   }
 
+  it should "perform chunking by title" taggedAs FastTest in {
+    val partitionOptions = Map(
+      "contentType" -> "text/html",
+      "titleFontSize" -> "14",
+      "chunkingStrategy" -> "byTitle",
+      "combineTextUnderNChars" -> "50")
+    val textDf = Partition(partitionOptions).partition(s"$htmlDirectory/fake-html.html")
+
+    val partitionDf = textDf.select(explode($"chunks.content"))
+    partitionDf.show(truncate = false)
+    assert(partitionDf.count() == 2)
+  }
+
 }
diff --git a/src/test/scala/com/johnsnowlabs/partition/TitleChunkerTest.scala b/src/test/scala/com/johnsnowlabs/partition/TitleChunkerTest.scala
