get spark to hoodie schema round trip conversion

Author: rahil-c

hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/io/storage/HoodieSparkParquetReader.java

@@ -57,6 +57,7 @@
 import org.apache.spark.sql.execution.datasources.parquet.SparkBasicSchemaEvolution;
 import org.apache.spark.sql.internal.SQLConf;
 import org.apache.spark.sql.sources.Filter;
+import org.apache.spark.sql.types.Metadata;
 import org.apache.spark.sql.types.StructType;
 
 import java.io.IOException;
@@ -204,7 +205,7 @@ public HoodieSchema getSchema() {
       MessageType messageType = getFileSchema();
       StructType structType = getStructSchema();
       schemaOption = Option.of(HoodieSparkSchemaConverters.toHoodieType(
-          structType, true, messageType.getName(), StringUtils.EMPTY_STRING));
+          structType, true, messageType.getName(), StringUtils.EMPTY_STRING, Metadata.empty()));
     }
     return schemaOption.get();
   }

hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/avro/HoodieSparkSchemaConverters.scala

@@ -229,10 +229,24 @@ object HoodieSparkSchemaConverters {
           val newRecordNames = existingRecordNames + fullName
           val fields = hoodieSchema.getFields.asScala.map { f =>
             val schemaType = toSqlTypeHelper(f.schema(), newRecordNames)
-            val metadata = if (f.doc().isPresent && !f.doc().get().isEmpty) {
-              new MetadataBuilder().putString("comment", f.doc().get()).build()
-            } else {
-              Metadata.empty
+            // Merge type-specific metadata (e.g., vector dimension) with field doc comment
+            val metadata = schemaType.metadata match {
+              case Some(typeMetadata) =>
+                // If there's type-specific metadata, use it as base
+                if (f.doc().isPresent && !f.doc().get().isEmpty) {
+                  // Merge doc comment into the metadata
+                  new MetadataBuilder().withMetadata(typeMetadata)
+                    .putString("comment", f.doc().get()).build()
+                } else {
+                  typeMetadata
+                }
+              case None =>
+                // No type-specific metadata, just use doc comment if present
+                if (f.doc().isPresent && !f.doc().get().isEmpty) {
+                  new MetadataBuilder().putString("comment", f.doc().get()).build()
+                } else {
+                  Metadata.empty
+                }
             }
             StructField(f.name(), schemaType.dataType, schemaType.nullable, metadata)
           }

hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/TestHoodieSchemaConversionUtils.scala

@@ -570,6 +570,119 @@ class TestHoodieSchemaConversionUtils extends FunSuite with Matchers {
     internalRowCompare(row1, row2, sparkSchema)
   }
 
+  test("test VECTOR type conversion - Spark to HoodieSchema") {
+    val metadata = new MetadataBuilder()
+      .putLong("hoodie.vector.dimension", 128)
+      .build()
+    val struct = new StructType()
+      .add("id", IntegerType, false)
+      .add("embedding", ArrayType(FloatType, containsNull = false), nullable = false, metadata)
+
+    val hoodieSchema = HoodieSchemaConversionUtils.convertStructTypeToHoodieSchema(
+      struct, "VectorTest", "test")
+
+    // Verify the embedding field is VECTOR type
+    val embeddingField = hoodieSchema.getField("embedding").get()
+    assert(embeddingField.schema().getType == HoodieSchemaType.VECTOR)
+    assert(embeddingField.schema().isInstanceOf[HoodieSchema.Vector])
+
+    val vectorSchema = embeddingField.schema().asInstanceOf[HoodieSchema.Vector]
+    assert(vectorSchema.getDimension == 128)
+    assert(vectorSchema.getVectorElementType == HoodieSchema.Vector.ELEMENT_TYPE_FLOAT)
+    assert(!embeddingField.isNullable())
+  }
+
+  test("test VECTOR type conversion - HoodieSchema to Spark") {
+    val vectorSchema = HoodieSchema.createVector(256, HoodieSchema.Vector.ELEMENT_TYPE_FLOAT)
+    val fields = java.util.Arrays.asList(
+      HoodieSchemaField.of("id", HoodieSchema.create(HoodieSchemaType.INT)),
+      HoodieSchemaField.of("embedding", vectorSchema)
+    )
+    val hoodieSchema = HoodieSchema.createRecord("VectorTest", "test", null, fields)
+
+    val structType = HoodieSchemaConversionUtils.convertHoodieSchemaToStructType(hoodieSchema)
+
+    // Verify the embedding field is ArrayType(FloatType)
+    assert(structType.fields.length == 2)
+    val embeddingField = structType.fields(1)
+    assert(embeddingField.name == "embedding")
+    assert(embeddingField.dataType.isInstanceOf[ArrayType])
+    assert(!embeddingField.nullable)
+
+    val arrayType = embeddingField.dataType.asInstanceOf[ArrayType]
+    assert(arrayType.elementType == FloatType)
+    assert(!arrayType.containsNull)
+
+    // Verify metadata contains dimension
+    assert(embeddingField.metadata.contains("hoodie.vector.dimension"))
+    assert(embeddingField.metadata.getLong("hoodie.vector.dimension") == 256)
+  }
+
+  test("test VECTOR round-trip conversion - Spark to HoodieSchema to Spark") {
+    val metadata = new MetadataBuilder()
+      .putLong("hoodie.vector.dimension", 512)
+      .build()
+    val originalStruct = new StructType()
+      .add("id", LongType, false)
+      .add("vector_field", ArrayType(FloatType, containsNull = false), nullable = false, metadata)
+      .add("name", StringType, true)
+
+    // Convert Spark -> HoodieSchema -> Spark
+    val hoodieSchema = HoodieSchemaConversionUtils.convertStructTypeToHoodieSchema(
+      originalStruct, "RoundTripTest", "test")
+    val convertedStruct = HoodieSchemaConversionUtils.convertHoodieSchemaToStructType(hoodieSchema)
+
+    // Verify structure is preserved
+    assert(convertedStruct.fields.length == originalStruct.fields.length)
+
+    // Verify vector field properties
+    val originalVectorField = originalStruct.fields(1)
+    val convertedVectorField = convertedStruct.fields(1)
+
+    assert(convertedVectorField.name == originalVectorField.name)
+    assert(convertedVectorField.dataType == originalVectorField.dataType)
+    assert(convertedVectorField.nullable == originalVectorField.nullable)
+
+    // Verify metadata is preserved
+    assert(convertedVectorField.metadata.contains("hoodie.vector.dimension"))
+    assert(convertedVectorField.metadata.getLong("hoodie.vector.dimension") == 512)
+
+    // Verify array properties
+    val convertedArrayType = convertedVectorField.dataType.asInstanceOf[ArrayType]
+    assert(convertedArrayType.elementType == FloatType)
+    assert(!convertedArrayType.containsNull)
+  }
+
+  test("test VECTOR round-trip conversion - HoodieSchema to Spark to HoodieSchema") {
+    val originalVectorSchema = HoodieSchema.createVector(1024, HoodieSchema.Vector.ELEMENT_TYPE_FLOAT)
+    val fields = java.util.Arrays.asList(
+      HoodieSchemaField.of("id", HoodieSchema.create(HoodieSchemaType.LONG)),
+      HoodieSchemaField.of("embedding", originalVectorSchema),
+      HoodieSchemaField.of("metadata", HoodieSchema.createNullable(HoodieSchemaType.STRING))
+    )
+    val originalHoodieSchema = HoodieSchema.createRecord("RoundTripTest", "test", null, fields)
+
+    // Convert HoodieSchema -> Spark -> HoodieSchema
+    val sparkStruct = HoodieSchemaConversionUtils.convertHoodieSchemaToStructType(originalHoodieSchema)
+    val convertedHoodieSchema = HoodieSchemaConversionUtils.convertStructTypeToHoodieSchema(
+      sparkStruct, "RoundTripTest", "test")
+
+    // Verify the vector field is preserved
+    val convertedEmbeddingField = convertedHoodieSchema.getField("embedding").get()
+    assert(convertedEmbeddingField.schema().getType == HoodieSchemaType.VECTOR)
+    assert(convertedEmbeddingField.schema().isInstanceOf[HoodieSchema.Vector])
+
+    val convertedVectorSchema = convertedEmbeddingField.schema().asInstanceOf[HoodieSchema.Vector]
+    assert(convertedVectorSchema.getDimension == 1024)
+    assert(convertedVectorSchema.getVectorElementType == HoodieSchema.Vector.ELEMENT_TYPE_FLOAT)
+    assert(!convertedEmbeddingField.isNullable())
+
+    // Verify other fields are preserved
+    assert(convertedHoodieSchema.getFields.size() == 3)
+    assert(convertedHoodieSchema.getField("id").get().schema().getType == HoodieSchemaType.LONG)
+    assert(convertedHoodieSchema.getField("metadata").get().isNullable())
+  }
+
   private def internalRowCompare(expected: Any, actual: Any, schema: DataType): Unit = {
     schema match {
       case StructType(fields) =>