[SPARK-54306] Annotate Variant columns with Variant logical type annotation

### What changes were proposed in this pull request?

This PR makes changes to the parquet writer to make it annotate variant columns with the parquet variant logical type annotation.

### Why are the changes needed?

The Parquet spec has formally adopted the Variant logical type, and therefore, Variant columns must be properly annotated in Spark 4.1.0 which depends on Parquet-java 1.16.0 which contains the variant logical type annotation.

This change is hidden behind a flag that is disabled by default until read support can be properly implemented.

### Does this PR introduce _any_ user-facing change?

Yes, Parquet files written by Spark 4.1.0 with the flag enabled (which it eventually will be by default) could contain the variant logical type annotation which readers without support for the type will not be able to read

### How was this patch tested?

Unit test to check if nested as well as top-level variants are properly annotated, and the data is being written correctly.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes #53005 from harshmotw-db/harshmotw-db/variant_annotation_write.

Authored-by: Harsh Motwani <harsh.motwani@databricks.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
(cherry picked from commit 5270c99)
Signed-off-by: Wenchen Fan <wenchen@databricks.com>

Author: harshmotw-db

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1585,6 +1585,14 @@ object SQLConf {
       .booleanConf
       .createWithDefault(true)
 
+  val PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE =
+    buildConf("spark.sql.parquet.variant.annotateLogicalType.enabled")
+      .doc("When enabled, Spark annotates the variant groups written to Parquet as the parquet " +
+        "variant logical type.")
+      .version("4.1.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val PARQUET_FIELD_ID_READ_ENABLED =
     buildConf("spark.sql.parquet.fieldId.read.enabled")
       .doc("Field ID is a native field of the Parquet schema spec. When enabled, Parquet readers " +
@@ -7638,6 +7646,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
 
   def parquetFieldIdWriteEnabled: Boolean = getConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED)
 
+  def parquetAnnotateVariantLogicalType: Boolean = getConf(PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE)
+
   def ignoreMissingParquetFieldId: Boolean = getConf(SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID)
 
   def legacyParquetNanosAsLong: Boolean = getConf(SQLConf.LEGACY_PARQUET_NANOS_AS_LONG)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala

@@ -373,6 +373,10 @@ class ParquetToSparkSchemaConverter(
 
     Option(field.getLogicalTypeAnnotation).fold(
       convertInternal(groupColumn, sparkReadType.map(_.asInstanceOf[StructType]))) {
+      // Temporary workaround to read Shredded variant data
+      case v: VariantLogicalTypeAnnotation if v.getSpecVersion == 1 && sparkReadType.isEmpty =>
+        convertInternal(groupColumn, None)
+
       // A Parquet list is represented as a 3-level structure:
       //
       //   <list-repetition> group <name> (LIST) {
@@ -552,7 +556,9 @@ class SparkToParquetSchemaConverter(
     writeLegacyParquetFormat: Boolean = SQLConf.PARQUET_WRITE_LEGACY_FORMAT.defaultValue.get,
     outputTimestampType: SQLConf.ParquetOutputTimestampType.Value =
       SQLConf.ParquetOutputTimestampType.INT96,
-    useFieldId: Boolean = SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.defaultValue.get) {
+    useFieldId: Boolean = SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.defaultValue.get,
+    annotateVariantLogicalType: Boolean =
+      SQLConf.PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE.defaultValue.get) {
 
   def this(conf: SQLConf) = this(
     writeLegacyParquetFormat = conf.writeLegacyParquetFormat,
@@ -563,7 +569,9 @@ class SparkToParquetSchemaConverter(
     writeLegacyParquetFormat = conf.get(SQLConf.PARQUET_WRITE_LEGACY_FORMAT.key).toBoolean,
     outputTimestampType = SQLConf.ParquetOutputTimestampType.withName(
       conf.get(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key)),
-    useFieldId = conf.get(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key).toBoolean)
+    useFieldId = conf.get(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key).toBoolean,
+    annotateVariantLogicalType =
+      conf.get(SQLConf.PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE.key).toBoolean)
 
   /**
    * Converts a Spark SQL [[StructType]] to a Parquet [[MessageType]].
@@ -817,14 +825,22 @@ class SparkToParquetSchemaConverter(
       // ===========
 
       case VariantType =>
-        Types.buildGroup(repetition)
+        (if (annotateVariantLogicalType) {
+          Types.buildGroup(repetition).as(LogicalTypeAnnotation.variantType(1))
+        } else {
+          Types.buildGroup(repetition)
+        })
           .addField(convertField(StructField("value", BinaryType, nullable = false), inShredded))
           .addField(convertField(StructField("metadata", BinaryType, nullable = false), inShredded))
           .named(field.name)
 
       case s: StructType if SparkShreddingUtils.isVariantShreddingStruct(s) =>
         // Variant struct takes a Variant and writes to Parquet as a shredded schema.
-        val group = Types.buildGroup(repetition)
+        val group = if (annotateVariantLogicalType) {
+          Types.buildGroup(repetition).as(LogicalTypeAnnotation.variantType(1))
+        } else {
+          Types.buildGroup(repetition)
+        }
         s.fields.foreach { f =>
           group.addField(convertField(f, inShredded = true))
         }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala

@@ -523,6 +523,10 @@ object ParquetUtils extends Logging {
       SQLConf.LEGACY_PARQUET_NANOS_AS_LONG.key,
       sqlConf.legacyParquetNanosAsLong.toString)
 
+    conf.set(
+      SQLConf.PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE.key,
+      sqlConf.parquetAnnotateVariantLogicalType.toString)
+
     // Sets compression scheme
     conf.set(ParquetOutputFormat.COMPRESSION, parquetOptions.compressionCodecClassName)
 

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVariantShreddingSuite.scala

@@ -19,11 +19,13 @@ package org.apache.spark.sql.execution.datasources.parquet
 
 import java.io.File
 
+import scala.jdk.CollectionConverters._
+
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.parquet.hadoop.ParquetFileReader
 import org.apache.parquet.hadoop.util.HadoopInputFile
-import org.apache.parquet.schema.{LogicalTypeAnnotation, PrimitiveType}
+import org.apache.parquet.schema.{LogicalTypeAnnotation, PrimitiveType, Type}
 import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
 
 import org.apache.spark.sql.{QueryTest, Row}
@@ -154,6 +156,73 @@ class ParquetVariantShreddingSuite extends QueryTest with ParquetTest with Share
     }
   }
 
+  test("variant logical type annotation") {
+    Seq(false, true).foreach { annotateVariantLogicalType =>
+      Seq(false, true).foreach { shredVariant =>
+        Seq(false, true).foreach { allowReadingShredded =>
+          withSQLConf(SQLConf.VARIANT_WRITE_SHREDDING_ENABLED.key -> shredVariant.toString,
+            SQLConf.VARIANT_INFER_SHREDDING_SCHEMA.key -> shredVariant.toString,
+            SQLConf.VARIANT_ALLOW_READING_SHREDDED.key ->
+              (allowReadingShredded || shredVariant).toString,
+            SQLConf.PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE.key ->
+              annotateVariantLogicalType.toString) {
+            def validateAnnotation(g: Type): Unit = {
+              if (annotateVariantLogicalType) {
+                assert(g.getLogicalTypeAnnotation == LogicalTypeAnnotation.variantType(1))
+              } else {
+                assert(g.getLogicalTypeAnnotation == null)
+              }
+            }
+            withTempDir { dir =>
+              // write parquet file
+              val df = spark.sql(
+                """
+                  | select
+                  |  id * 2 i,
+                  |  to_variant_object(named_struct('id', id)) v,
+                  |  named_struct('i', (id * 2)::string,
+                  |     'nv', to_variant_object(named_struct('id', 30 + id))) ns,
+                  |  array(to_variant_object(named_struct('id', 10 + id))) av,
+                  |  map('v2', to_variant_object(named_struct('id', 20 + id))) mv
+                  |  from range(0,3,1,1)""".stripMargin)
+              df.write.mode("overwrite").parquet(dir.getAbsolutePath)
+              val file = dir.listFiles().find(_.getName.endsWith(".parquet")).get
+              val parquetFilePath = file.getAbsolutePath
+              val inputFile = HadoopInputFile.fromPath(new Path(parquetFilePath),
+                new Configuration())
+              val reader = ParquetFileReader.open(inputFile)
+              val footer = reader.getFooter
+              val schema = footer.getFileMetaData.getSchema
+              val vGroup = schema.getType(schema.getFieldIndex("v"))
+              validateAnnotation(vGroup)
+              assert(vGroup.asGroupType().getFields.asScala.toSeq
+                .exists(_.getName == "typed_value") == shredVariant)
+              val nsGroup = schema.getType(schema.getFieldIndex("ns")).asGroupType()
+              val nvGroup = nsGroup.getType(nsGroup.getFieldIndex("nv"))
+              validateAnnotation(nvGroup)
+              val avGroup = schema.getType(schema.getFieldIndex("av")).asGroupType()
+              val avList = avGroup.getType(avGroup.getFieldIndex("list")).asGroupType()
+              val avElement = avList.getType(avList.getFieldIndex("element"))
+              validateAnnotation(avElement)
+              val mvGroup = schema.getType(schema.getFieldIndex("mv")).asGroupType()
+              val mvList = mvGroup.getType(mvGroup.getFieldIndex("key_value")).asGroupType()
+              val mvValue = mvList.getType(mvList.getFieldIndex("value"))
+              validateAnnotation(mvValue)
+              // verify result
+              val result = spark.read.format("parquet")
+                .schema("v variant, ns struct<nv variant>, av array<variant>, " +
+                  "mv map<string, variant>")
+                .load(dir.getAbsolutePath)
+                .selectExpr("v:id::int i1", "ns.nv:id::int i2", "av[0]:id::int i3",
+                  "mv['v2']:id::int i4")
+              checkAnswer(result, Array(Row(0, 30, 10, 20), Row(1, 31, 11, 21), Row(2, 32, 12, 22)))
+              reader.close()
+            }
+          }
+        }
+      }
+    }
+  }
 
   testWithTempDir("write shredded variant basic") { dir =>
     val schema = "a int, b string, c decimal(15, 1)"