Extend metadata and metadatabuilder

Signed-off-by: Peng Huo <[email protected]>

Author: penghuo

Diff for: flint-spark-integration/src/main/scala/org/apache/spark/sql/flint/datatype/FlintDataType.scala

@@ -12,6 +12,7 @@ import org.json4s.jackson.JsonMethods
 import org.json4s.native.Serialization
 
 import org.apache.spark.sql.catalyst.util.DateFormatter
+import org.apache.spark.sql.flint.datatype.FlintMetadataExtensions.{MetadataBuilderExtension, MetadataExtension}
 import org.apache.spark.sql.types._
 
 /**
@@ -97,14 +98,12 @@ object FlintDataType {
 
       // Text with possible multi-fields
       case JString("text") =>
-        FlintMetadataHelper.addTextFieldMetadata(metadataBuilder)
+        metadataBuilder.withTextField()
         (fieldProperties \ "fields") match {
           case fields: JObject =>
-            FlintMetadataHelper.addMultiFieldMetadata(
-              metadataBuilder,
-              fields.obj.map { case (name, props) =>
-                (s"$fieldName.$name", (props \ "type").extract[String])
-              }.toMap)
+            metadataBuilder.withMultiFields(fields.obj.map { case (name, props) =>
+              (s"$fieldName.$name", (props \ "type").extract[String])
+            }.toMap)
             StringType
           case _ => StringType
         }
@@ -164,7 +163,7 @@ object FlintDataType {
 
       // string
       case StringType | _: VarcharType | _: CharType =>
-        if (FlintMetadataHelper.isTextField(metadata)) {
+        if (metadata.isTextField) {
           JObject("type" -> JString("text"))
         } else {
           JObject("type" -> JString("keyword"))

Diff for: flint-spark-integration/src/main/scala/org/apache/spark/sql/flint/datatype/FlintMetadataExtensions.scala

@@ -0,0 +1,90 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.apache.spark.sql.flint.datatype
+
+import org.apache.spark.sql.types.{MetadataBuilder, _}
+
+/**
+ * Helper class for handling Flint metadata operations
+ */
+object FlintMetadataExtensions {
+  // OpenSearch Mappings. https://opensearch.org/docs/latest/field-types/
+  val OS_TYPE_KEY = "osType"
+  val FIELDS_NAMES_KEY = "fields"
+
+  // OpenSearch field types. https://opensearch.org/docs/latest/field-types/supported-field-types/index/
+  val TEXT_TYPE = "text"
+  val KEYWORD_TYPE = "keyword"
+
+  implicit class MetadataExtension(val metadata: Metadata) {
+
+    /**
+     * Determines if this metadata represents a text field.
+     *
+     * @return
+     *   `true` if the metadata contains an `osType` property set to "text", `false` otherwise
+     */
+    def isTextField: Boolean =
+      metadata.contains(OS_TYPE_KEY) && metadata.getString(OS_TYPE_KEY) == TEXT_TYPE
+
+    /**
+     * Retrieves the first keyword subfield from multi-field definitions.
+     *
+     * @return
+     *   [[Some]] containing the first keyword subfield name if defined, [[None]] if no keyword
+     *   subfields exist
+     */
+    def keywordSubfield: Option[String] = {
+      if (metadata.contains(FIELDS_NAMES_KEY)) {
+        val multiFieldMetadata = metadata.getMetadata(FIELDS_NAMES_KEY)
+        multiFieldMetadata
+          .getStringArray(KEYWORD_TYPE)
+          .headOption
+      } else {
+        None
+      }
+    }
+  }
+
+  implicit class MetadataBuilderExtension(val builder: MetadataBuilder) {
+
+    /**
+     * Marks the field as a text field in OpenSearch metadata.
+     *
+     * @return
+     *   the modified [[MetadataBuilder]] instance for method chaining
+     * @see
+     *   [[https://opensearch.org/docs/latest/field-types/supported-field-types/text/ Text Field Documentation]]
+     */
+    def withTextField(): MetadataBuilder = {
+      builder.putString(OS_TYPE_KEY, TEXT_TYPE)
+      builder
+    }
+
+    /**
+     * Adds multi-field definitions to the metadata.
+     *
+     * @param fields
+     *   Map where keys are field names and values are OpenSearch field types
+     * @return
+     *   the modified [[MetadataBuilder]] instance for method chaining
+     * @example
+     *   {{{builder.withMultiFields(Map( "raw" -> KEYWORD_TYPE, "analyzed" -> TEXT_TYPE ))}}}
+     * @note
+     *   Groups fields by type and stores them under the `fields` metadata key
+     */
+    def withMultiFields(fields: Map[String, String]): MetadataBuilder = {
+      val nestedBuilder = new MetadataBuilder()
+      fields
+        .groupBy { case (_, fieldType) => fieldType }
+        .foreach { case (fieldType, entries) =>
+          nestedBuilder.putStringArray(fieldType, entries.keys.toArray)
+        }
+      builder.putMetadata(FIELDS_NAMES_KEY, nestedBuilder.build())
+      builder
+    }
+  }
+}

Diff for: flint-spark-integration/src/main/scala/org/apache/spark/sql/flint/datatype/FlintMetadataHelper.scala

@@ -13,7 +13,8 @@ import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.parseColumnPath
 import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, LiteralValue}
 import org.apache.spark.sql.connector.expressions.filter.{And, Predicate}
 import org.apache.spark.sql.flint.datatype.FlintDataType.STRICT_DATE_OPTIONAL_TIME_FORMATTER_WITH_NANOS
-import org.apache.spark.sql.flint.datatype.FlintMetadataHelper
+import org.apache.spark.sql.flint.datatype.FlintMetadataExtensions
+import org.apache.spark.sql.flint.datatype.FlintMetadataExtensions.MetadataExtension
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
@@ -157,7 +158,7 @@ case class FlintQueryCompiler(schema: StructType) {
       case Some((_, field)) =>
         field.dataType match {
           case StringType =>
-            FlintMetadataHelper.isTextField(field.metadata)
+            field.metadata.isTextField
           case _ => false
         }
       case None => false
@@ -169,8 +170,7 @@ case class FlintQueryCompiler(schema: StructType) {
    */
   protected def getKeywordSubfield(attribute: String): Option[String] = {
     schema.apply(attribute) match {
-      case StructField(_, StringType, _, metadata) =>
-        FlintMetadataHelper.getKeywordSubfield(metadata)
+      case StructField(_, StringType, _, metadata) => metadata.keywordSubfield
       case _ => None
     }
   }

Diff for: flint-spark-integration/src/main/scala/org/apache/spark/sql/flint/storage/FlintQueryCompiler.scala

@@ -10,6 +10,7 @@ import org.json4s.jackson.JsonMethods
 import org.scalatest.matchers.should.Matchers
 
 import org.apache.spark.FlintSuite
+import org.apache.spark.sql.flint.datatype.FlintMetadataExtensions.MetadataBuilderExtension
 import org.apache.spark.sql.types._
 
 class FlintDataTypeSuite extends FlintSuite with Matchers {
@@ -221,12 +222,9 @@ class FlintDataTypeSuite extends FlintSuite with Matchers {
                           |    }
                           |  }
                           |}""".stripMargin
-    val mb = new MetadataBuilder()
-    FlintMetadataHelper.addTextFieldMetadata(mb)
-    FlintMetadataHelper.addMultiFieldMetadata(
-      mb,
-      Map("city.raw" -> "keyword", "city.keyword" -> "keyword"))
-    val metadata = mb.build()
+    val metadata = new MetadataBuilder().withTextField
+      .withMultiFields(Map("city.raw" -> "keyword", "city.keyword" -> "keyword"))
+      .build()
     val expectedStructType = StructType(StructField("city", StringType, true, metadata) :: Nil)
 
     FlintDataType.deserialize(flintDataType) should contain theSameElementsAs expectedStructType
@@ -241,9 +239,7 @@ class FlintDataTypeSuite extends FlintSuite with Matchers {
                           |  }
                           |}""".stripMargin
 
-    val mb = new MetadataBuilder()
-    FlintMetadataHelper.addTextFieldMetadata(mb)
-    val metadata = mb.build()
+    val metadata = new MetadataBuilder().withTextField().build()
 
     val expectedStructType =
       StructType(StructField("description", StringType, true, metadata) :: Nil)

Diff for: flint-spark-integration/src/test/scala/org/apache/spark/sql/flint/datatype/FlintDataTypeSuite.scala

@@ -8,79 +8,83 @@ package org.apache.spark.sql.flint.datatype
 import org.scalatest.matchers.should.Matchers
 
 import org.apache.spark.FlintSuite
+import org.apache.spark.sql.flint.datatype.FlintMetadataExtensions.{MetadataBuilderExtension, MetadataExtension}
 import org.apache.spark.sql.types._
 
-class FlintMetadataHelperSuite extends FlintSuite with Matchers {
+class FlintMetadataExtensionsSuite extends FlintSuite with Matchers {
 
   test("isTextField returns true when osType is text") {
     val builder = new MetadataBuilder()
-      .putString(FlintMetadataHelper.OS_TYPE_KEY, FlintMetadataHelper.TEXT_TYPE)
+      .putString(FlintMetadataExtensions.OS_TYPE_KEY, FlintMetadataExtensions.TEXT_TYPE)
     val metadata: Metadata = builder.build()
-    assert(FlintMetadataHelper.isTextField(metadata))
+    assert(metadata.isTextField)
   }
 
   test("isTextField returns false when osType is not text") {
-    val builder = new MetadataBuilder().putString(FlintMetadataHelper.OS_TYPE_KEY, "non-text")
+    val builder = new MetadataBuilder().putString(FlintMetadataExtensions.OS_TYPE_KEY, "non-text")
     val metadata: Metadata = builder.build()
-    assert(!FlintMetadataHelper.isTextField(metadata))
+    assert(!metadata.isTextField)
   }
 
   test("addTextFieldMetadata sets osType to text") {
     val builder = new MetadataBuilder()
-    val updatedBuilder = FlintMetadataHelper.addTextFieldMetadata(builder)
+    val updatedBuilder = builder.withTextField
     val metadata: Metadata = updatedBuilder.build()
-    assert(metadata.getString(FlintMetadataHelper.OS_TYPE_KEY) == FlintMetadataHelper.TEXT_TYPE)
+    assert(
+      metadata.getString(
+        FlintMetadataExtensions.OS_TYPE_KEY) == FlintMetadataExtensions.TEXT_TYPE)
   }
 
   test("addMultiFieldMetadata groups fields by field type") {
     val builder = new MetadataBuilder()
     val fields = Map(
-      "field1" -> FlintMetadataHelper.TEXT_TYPE,
-      "field2" -> FlintMetadataHelper.KEYWORD_TYPE,
-      "field3" -> FlintMetadataHelper.KEYWORD_TYPE)
-    val updatedBuilder = FlintMetadataHelper.addMultiFieldMetadata(builder, fields)
+      "field1" -> FlintMetadataExtensions.TEXT_TYPE,
+      "field2" -> FlintMetadataExtensions.KEYWORD_TYPE,
+      "field3" -> FlintMetadataExtensions.KEYWORD_TYPE)
+    val updatedBuilder = builder.withMultiFields(fields)
     val metadata: Metadata = updatedBuilder.build()
 
     // Verify that multi-field metadata is added under FIELDS_NAMES_KEY.
-    assert(metadata.contains(FlintMetadataHelper.FIELDS_NAMES_KEY))
-    val multiFieldMetadata: Metadata = metadata.getMetadata(FlintMetadataHelper.FIELDS_NAMES_KEY)
+    assert(metadata.contains(FlintMetadataExtensions.FIELDS_NAMES_KEY))
+    val multiFieldMetadata: Metadata =
+      metadata.getMetadata(FlintMetadataExtensions.FIELDS_NAMES_KEY)
 
     // Verify text type field grouping.
-    assert(multiFieldMetadata.contains(FlintMetadataHelper.TEXT_TYPE))
-    val textFields = multiFieldMetadata.getStringArray(FlintMetadataHelper.TEXT_TYPE)
+    assert(multiFieldMetadata.contains(FlintMetadataExtensions.TEXT_TYPE))
+    val textFields = multiFieldMetadata.getStringArray(FlintMetadataExtensions.TEXT_TYPE)
     assert(textFields.sameElements(Array("field1")))
 
     // Verify keyword type field grouping.
-    assert(multiFieldMetadata.contains(FlintMetadataHelper.KEYWORD_TYPE))
-    val keywordFields = multiFieldMetadata.getStringArray(FlintMetadataHelper.KEYWORD_TYPE)
+    assert(multiFieldMetadata.contains(FlintMetadataExtensions.KEYWORD_TYPE))
+    val keywordFields = multiFieldMetadata.getStringArray(FlintMetadataExtensions.KEYWORD_TYPE)
     // Since the order of grouping may vary, compare sorted arrays.
     assert(keywordFields.sorted.sameElements(Array("field2", "field3")))
   }
 
   test("getKeywordSubfield returns the first keyword field if available") {
     val builder = new MetadataBuilder()
     val fields = Map(
-      "field1" -> FlintMetadataHelper.TEXT_TYPE,
-      "field2" -> FlintMetadataHelper.KEYWORD_TYPE,
-      "field3" -> FlintMetadataHelper.KEYWORD_TYPE)
-    val updatedBuilder = FlintMetadataHelper.addMultiFieldMetadata(builder, fields)
+      "field1" -> FlintMetadataExtensions.TEXT_TYPE,
+      "field2" -> FlintMetadataExtensions.KEYWORD_TYPE,
+      "field3" -> FlintMetadataExtensions.KEYWORD_TYPE)
+    val updatedBuilder = builder.withMultiFields(fields)
     val metadata: Metadata = updatedBuilder.build()
 
     // Retrieve keyword fields from the nested metadata.
-    val multiFieldMetadata = metadata.getMetadata(FlintMetadataHelper.FIELDS_NAMES_KEY)
-    val keywordFields = multiFieldMetadata.getStringArray(FlintMetadataHelper.KEYWORD_TYPE)
+    val multiFieldMetadata = metadata.getMetadata(FlintMetadataExtensions.FIELDS_NAMES_KEY)
+    val keywordFields = multiFieldMetadata.getStringArray(FlintMetadataExtensions.KEYWORD_TYPE)
 
     // Expect the first keyword field.
-    assert(FlintMetadataHelper.getKeywordSubfield(metadata) == keywordFields.headOption)
+    assert(metadata.keywordSubfield == keywordFields.headOption)
   }
 
   test("getKeywordSubfield returns None if no keyword field exists") {
     val builder = new MetadataBuilder()
-    val fields = Map("field1" -> FlintMetadataHelper.TEXT_TYPE)
-    val updatedBuilder = FlintMetadataHelper.addMultiFieldMetadata(builder, fields)
+    val fields = Map("field1" -> FlintMetadataExtensions.TEXT_TYPE)
+    val updatedBuilder = builder.withMultiFields(fields)
     val metadata: Metadata = updatedBuilder.build()
 
     // Since there is no keyword type, getKeywordSubfield should return None.
-    assert(FlintMetadataHelper.getKeywordSubfield(metadata).isEmpty)
+    assert(metadata.keywordSubfield.isEmpty)
   }
 }