[SPARKNLP-1299] Add Hierarchical Element Identification to HTMLReader

Author: danilojsl

src/main/scala/com/johnsnowlabs/reader/HTMLReader.scala

@@ -28,6 +28,9 @@ import org.jsoup.nodes.{Document, Element, Node, TextNode}
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
+import java.security.MessageDigest
+import java.util.UUID
+
 
 /** Class to parse and read HTML files.
   *
@@ -164,8 +167,8 @@ class HTMLReader(
   private def startTraversalFromBody(document: Document): Array[HTMLElement] = {
     try {
       val body = document.body()
-      val elements = extractElements(body)
       val docTitle = document.title().trim
+      val elements = extractElements(body)
 
       if (docTitle.nonEmpty && includeTitleTag) {
         val titleElem = HTMLElement(
@@ -205,6 +208,11 @@ class HTMLReader(
     val trackingNodes = mutable.Map[Node, NodeMetadata]()
     var pageNumber = 1
 
+    // Track parent-child hierarchy
+    var currentParentId: Option[String] = None
+
+    def newUUID(): String = UUID.randomUUID().toString
+
     def isNodeHidden(node: Node): Boolean = {
       node match {
         case elem: Element =>
@@ -237,15 +245,13 @@ class HTMLReader(
               trackingNodes(elem).visited = true
               val text = elem.ownText().trim
               if (text.nonEmpty) textBuffer += text
-              // Recursively collect text from all child nodes
               elem.childNodes().asScala.foreach(traverseAndCollect)
 
             case _ => // Ignore other node types
           }
         }
       }
 
-      // Start traversal for each node in the list
       nodes.foreach(traverseAndCollect)
       textBuffer.mkString(" ").replaceAll("\\s+", " ").trim
     }
@@ -261,9 +267,7 @@ class HTMLReader(
           NodeMetadata(tagName = tagName, hidden = isNodeHidden(childNode), visited = false))
       }
 
-      if (trackingNodes(node).hidden) {
-        return
-      }
+      if (trackingNodes(node).hidden) return
 
       node match {
         case element: Element =>
@@ -279,48 +283,53 @@ class HTMLReader(
               val linkText = element.text().trim
               if (href.nonEmpty && linkText.nonEmpty && !visitedNode) {
                 trackingNodes(element).visited = true
+                pageMetadata("element_id") = newUUID()
+                currentParentId.foreach(pid => pageMetadata("parent_id") = pid)
                 elements += HTMLElement(
                   ElementType.LINK,
                   content = s"[$linkText]($href)",
                   metadata = pageMetadata)
               }
+
             case "table" =>
               pageMetadata("sentence") = sentenceIndex.toString
               sentenceIndex += 1
               val tableContent = outputFormat match {
-                case "plain-text" =>
-                  extractNestedTableContent(element).trim
+                case "plain-text" => extractNestedTableContent(element).trim
                 case "html-table" =>
                   element
                     .outerHtml()
                     .replaceAll("\\n", "")
                     .replaceAll(">\\s+<", "><")
                     .replaceAll("^\\s+|\\s+$", "")
-                case "json-table" =>
-                  tableElementToJson(element)
-                case _ =>
-                  extractNestedTableContent(element).trim
+                case "json-table" => tableElementToJson(element)
+                case _ => extractNestedTableContent(element).trim
               }
               if (tableContent.nonEmpty && !visitedNode) {
                 trackingNodes(element).visited = true
+                pageMetadata("element_id") = newUUID()
+                currentParentId.foreach(pid => pageMetadata("parent_id") = pid)
                 elements += HTMLElement(
                   ElementType.TABLE,
                   content = tableContent,
                   metadata = pageMetadata)
               }
+
             case "li" =>
               pageMetadata("sentence") = sentenceIndex.toString
               sentenceIndex += 1
               val itemText = element.text().trim
               if (itemText.nonEmpty && !visitedNode) {
                 trackingNodes(element).visited = true
+                pageMetadata("element_id") = newUUID()
+                currentParentId.foreach(pid => pageMetadata("parent_id") = pid)
                 elements += HTMLElement(
                   ElementType.LIST_ITEM,
                   content = itemText,
                   metadata = pageMetadata)
               }
+
             case "pre" =>
-              // A <pre> tag typically contains a <code> child
               val codeElem = element.getElementsByTag("code").first()
               val codeText =
                 if (codeElem != null) codeElem.text().trim
@@ -329,22 +338,22 @@ class HTMLReader(
                 pageMetadata("sentence") = sentenceIndex.toString
                 sentenceIndex += 1
                 trackingNodes(element).visited = true
+                pageMetadata("element_id") = newUUID()
+                currentParentId.foreach(pid => pageMetadata("parent_id") = pid)
                 elements += HTMLElement(
                   ElementType.UNCATEGORIZED_TEXT,
                   content = codeText,
                   metadata = pageMetadata)
               }
+
             case tag if isParagraphLikeElement(element) =>
               if (!visitedNode) {
                 val classType = classifyParagraphElement(element)
-
-                // Traverse children first so that <img>, <a>, etc. inside the paragraph are processed
                 element.childNodes().asScala.foreach { childNode =>
                   val tagName = getTagName(childNode)
                   traverse(childNode, tagName)
                 }
 
-                // Now handle the paragraph itself
                 classType match {
                   case ElementType.NARRATIVE_TEXT =>
                     val childNodes = element.childNodes().asScala.toList
@@ -353,6 +362,8 @@ class HTMLReader(
                       pageMetadata("sentence") = sentenceIndex.toString
                       sentenceIndex += 1
                       trackingNodes(element).visited = true
+                      pageMetadata("element_id") = newUUID()
+                      currentParentId.foreach(pid => pageMetadata("parent_id") = pid)
                       elements += HTMLElement(
                         ElementType.NARRATIVE_TEXT,
                         content = aggregatedText,
@@ -365,10 +376,13 @@ class HTMLReader(
                       pageMetadata("sentence") = sentenceIndex.toString
                       sentenceIndex += 1
                       trackingNodes(element).visited = true
+                      val titleId = newUUID()
+                      pageMetadata("element_id") = titleId
                       elements += HTMLElement(
                         ElementType.TITLE,
                         content = titleText,
                         metadata = pageMetadata)
+                      currentParentId = Some(titleId)
                     }
 
                   case ElementType.UNCATEGORIZED_TEXT =>
@@ -377,28 +391,36 @@ class HTMLReader(
                       pageMetadata("sentence") = sentenceIndex.toString
                       sentenceIndex += 1
                       trackingNodes(element).visited = true
+                      pageMetadata("element_id") = newUUID()
+                      currentParentId.foreach(pid => pageMetadata("parent_id") = pid)
                       elements += HTMLElement(
                         ElementType.UNCATEGORIZED_TEXT,
                         content = text,
                         metadata = pageMetadata)
                     }
                 }
               }
+
             case _ if isTitleElement(element) && !visitedNode =>
               trackingNodes(element).visited = true
               val titleText = element.text().trim
               if (titleText.nonEmpty) {
                 pageMetadata("sentence") = sentenceIndex.toString
                 sentenceIndex += 1
+                val titleId = newUUID()
+                pageMetadata("element_id") = titleId
                 elements += HTMLElement(
                   ElementType.TITLE,
                   content = titleText,
                   metadata = pageMetadata)
+                currentParentId = Some(titleId)
               }
+
             case "hr" =>
               if (element.attr("style").toLowerCase.contains("page-break")) {
                 pageNumber = pageNumber + 1
               }
+
             case "img" =>
               pageMetadata("sentence") = sentenceIndex.toString
               sentenceIndex += 1
@@ -411,7 +433,6 @@ class HTMLReader(
                 val height = element.attr("height").trim
 
                 val imgMetadata = mutable.Map[String, String]("alt" -> alt) ++ pageMetadata
-
                 var contentValue = src
                 if (isBase64) {
                   val commaIndex = src.indexOf(',')
@@ -422,14 +443,16 @@ class HTMLReader(
                     contentValue = base64Payload
                   }
                 }
-
                 if (width.nonEmpty) imgMetadata("width") = width
                 if (height.nonEmpty) imgMetadata("height") = height
+                imgMetadata("element_id") = newUUID()
+                currentParentId.foreach(pid => imgMetadata("parent_id") = pid)
                 elements += HTMLElement(
                   ElementType.IMAGE,
                   content = contentValue,
                   metadata = imgMetadata)
               }
+
             case _ =>
               element.childNodes().asScala.foreach { childNode =>
                 val tagName = getTagName(childNode)
@@ -440,7 +463,6 @@ class HTMLReader(
       }
     }
 
-    // Start traversal from the root node
     val tagName = getTagName(root)
     traverse(root, tagName)
     elements.toArray

src/test/resources/reader/html/simple-book.html

@@ -0,0 +1,64 @@
+<!-- File: /index.html -->
+<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8" />
+  <title>Simple Book: 3 Chapters</title>
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <style>
+    /* Keep navigation readable on small screens */
+    body { font-family: system-ui, -apple-system, Segoe UI, Roboto, Arial, sans-serif; line-height: 1.6; margin: 2rem; }
+    nav ul { list-style: none; padding: 0; }
+    nav li { margin: 0.25rem 0; }
+    a { text-decoration: none; }
+    a:hover { text-decoration: underline; }
+    hr { margin: 2rem 0; }
+    .back { display: inline-block; margin-top: 0.5rem; }
+  </style>
+</head>
+<body>
+  <h1 id="top">Simple Book</h1>
+
+  <nav aria-label="Chapter index">
+    <h2>Index</h2>
+    <ul>
+      <li><a href="#chapter-1">Chapter 1: Beginnings</a></li>
+      <li><a href="#chapter-2">Chapter 2: Middle Path</a></li>
+      <li><a href="#chapter-3">Chapter 3: Finishing Touch</a></li>
+    </ul>
+  </nav>
+
+  <hr />
+
+  <section id="chapter-1">
+    <h2>Chapter 1: Beginnings</h2>
+    <p>
+      Every project starts with a simple idea and a clear intention. In this chapter, we set the stage and outline the basic goals. 
+      Small steps help build momentum and reduce uncertainty. With a plan in place, moving forward becomes much easier.
+    </p>
+    <a class="back" href="#top">Back to top</a>
+  </section>
+
+  <hr />
+
+  <section id="chapter-2">
+    <h2>Chapter 2: Middle Path</h2>
+    <p>
+      Progress is rarely a straight line, and that is perfectly fine. Here we adjust our approach based on what we learn. 
+      Iteration helps refine ideas and improves the final outcome. Staying flexible keeps the project healthy and on track.
+    </p>
+    <a class="back" href="#top">Back to top</a>
+  </section>
+
+  <hr />
+
+  <section id="chapter-3">
+    <h2>Chapter 3: Finishing Touch</h2>
+    <p>
+      The final phase focuses on clarity and polish. We review the work, remove distractions, and keep what matters. 
+      A simple, tidy result is easier to use and maintain. With that, the project is ready to share.
+    </p>
+    <a class="back" href="#top">Back to top</a>
+  </section>
+</body>
+</html>

src/test/scala/com/johnsnowlabs/reader/HTMLReaderTest.scala

@@ -125,6 +125,7 @@ class HTMLReaderTest extends AnyFlatSpec {
   it should "correctly parse caption and th tags" taggedAs FastTest in {
     val HTMLReader = new HTMLReader()
     val htmlDF = HTMLReader.read(s"$htmlFilesDirectory/example-caption-th.html")
+    htmlDF.show(truncate = false)
     val titleDF = htmlDF
       .select(explode(col("html")).as("exploded_html"))
       .filter(col("exploded_html.elementType") === ElementType.TABLE)
@@ -185,4 +186,77 @@ class HTMLReaderTest extends AnyFlatSpec {
     assert(imagesDF.count() == 1)
   }
 
+  it should "include parent and element ids" taggedAs FastTest in {
+    val HTMLReader = new HTMLReader()
+    val htmlDF = HTMLReader.read(s"$htmlFilesDirectory/simple-book.html")
+    htmlDF.show(truncate = false)
+    val parentChildDF = htmlDF
+      .select(explode(col("html")).as("exploded_html"))
+
+    parentChildDF.show(truncate = false)
+
+//    assert(parentChildDF.count() == 3)
+  }
+
+  it should "produce valid element_id and parent_id relationships" taggedAs FastTest in {
+    val HTMLReader = new HTMLReader()
+    val htmlDF = HTMLReader.read(s"$htmlFilesDirectory/simple-book.html")
+
+    val explodedDF = htmlDF
+      .select(explode(col("html")).as("elem"))
+      .select(
+        col("elem.elementType").as("elementType"),
+        col("elem.content").as("content"),
+        col("elem.metadata").as("metadata")
+      )
+      .withColumn("element_id", col("metadata")("element_id"))
+      .withColumn("parent_id", col("metadata")("parent_id"))
+      .cache() // << important to prevent recomputation inconsistencies
+
+    val allElementIds = explodedDF
+      .select("element_id")
+      .where(col("element_id").isNotNull)
+      .distinct()
+      .collect()
+      .map(_.getString(0))
+      .toSet
+
+    val allParentIds = explodedDF
+      .select("parent_id")
+      .where(col("parent_id").isNotNull)
+      .distinct()
+      .collect()
+      .map(_.getString(0))
+      .toSet
+
+    // 1. There should be at least one element with an element_id
+    assert(allElementIds.nonEmpty, "No elements have element_id metadata")
+
+    // 2. There should be at least one element with a parent_id
+    assert(allParentIds.nonEmpty, "No elements have parent_id metadata")
+
+    // 3. Every parent_id should exist as an element_id
+    val missingParents = allParentIds.diff(allElementIds)
+    assert(
+      missingParents.isEmpty,
+      s"Some parent_ids do not correspond to existing element_ids: $missingParents"
+    )
+
+    // 4. Each parent should have at least one child
+    val parentChildCount = explodedDF
+      .filter(col("parent_id").isNotNull)
+      .groupBy("parent_id")
+      .count()
+      .collect()
+      .map(r => r.getString(0) -> r.getLong(1))
+      .toMap
+
+    assert(
+      parentChildCount.nonEmpty && parentChildCount.values.forall(_ >= 1),
+      "Each parent_id should have at least one child element"
+    )
+  }
+
+
+
 }