[SPARK-58949][CORE] Use unsigned UTF-8 byte ordering for Variant object fields

common/variant/src/main/java/org/apache/spark/types/variant/Variant.java

@@ -144,18 +144,17 @@ public int objectSize() {
   // It is only legal to call it when `getType()` is `Type.OBJECT`.
   public Variant getFieldByKey(String key) {
     return handleObject(value, pos, (size, idSize, offsetSize, idStart, offsetStart, dataStart) -> {
-      // Use linear search for a short list. Switch to binary search when the length reaches
-      // `BINARY_SEARCH_THRESHOLD`.
-      final int BINARY_SEARCH_THRESHOLD = 32;
-      if (size < BINARY_SEARCH_THRESHOLD) {
-        for (int i = 0; i < size; ++i) {
-          int id = readUnsigned(value, idStart + idSize * i, idSize);
-          if (key.equals(getMetadataKey(metadata, id))) {
-            int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
-            return new Variant(value, metadata, dataStart + offset);
-          }
+      byte[] keyBytes = encodeKey(key);
+      int numAttempts = 1;
+      // UTF-8 and UTF-16 order can differ only for keys with a code unit at or above U+D800.
+      for (int i = 0; i < key.length(); ++i) {
+        if (key.charAt(i) >= Character.MIN_SURROGATE) {
+          numAttempts = 2;
+          break;
         }
-      } else {
+      }
+      // Search the spec's UTF-8 order first, then the UTF-16 order written by older Spark versions.
+      for (int attempt = 0; attempt < numAttempts; ++attempt) {
         int low = 0;
         int high = size - 1;
         while (low <= high) {
@@ -164,7 +163,10 @@ public Variant getFieldByKey(String key) {
           // overflows int.
           int mid = (low + high) >>> 1;
           int id = readUnsigned(value, idStart + idSize * mid, idSize);
-          int cmp = getMetadataKey(metadata, id).compareTo(key);
+          String midKey = getMetadataKey(metadata, id);
+          int cmp = attempt == 0
+              ? compareKeys(encodeKey(midKey), keyBytes)
+              : midKey.compareTo(key);
           if (cmp < 0) {
             low = mid + 1;
           } else if (cmp > 0) {

common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java

@@ -376,7 +376,7 @@ public int addKey(String key) {
     } else {
       id = dictionaryKeys.size();
       dictionary.put(key, id);
-      dictionaryKeys.add(key.getBytes(StandardCharsets.UTF_8));
+      dictionaryKeys.add(encodeKey(key));
     }
     return id;
   }
@@ -994,6 +994,7 @@ public static final class FieldEntry implements Comparable<FieldEntry> {
     final String key;
     final int id;
     final int offset;
+    private byte[] keyBytes;
 
     public FieldEntry(String key, int id, int offset) {
       this.key = key;
@@ -1005,9 +1006,16 @@ FieldEntry withNewOffset(int newOffset) {
       return new FieldEntry(key, id, newOffset);
     }
 
+    private byte[] keyBytes() {
+      if (keyBytes == null) {
+        keyBytes = encodeKey(key);
+      }
+      return keyBytes;
+    }
+
     @Override
     public int compareTo(FieldEntry other) {
-      return key.compareTo(other.key);
+      return compareKeys(keyBytes(), other.keyBytes());
     }
   }
 

common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java

@@ -688,6 +688,16 @@ private static void validateImpl(byte[] value, byte[] metadata, int pos) {
     }
   }
 
+  // Encode an object field key for comparison in the order required by the Variant spec.
+  public static byte[] encodeKey(String key) {
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  // Compare UTF-8-encoded object field keys using unsigned lexicographic byte ordering.
+  public static int compareKeys(byte[] left, byte[] right) {
+    return Arrays.compareUnsigned(left, right);
+  }
+
   // Get a key at `id` in the variant metadata.
   // Throw `MALFORMED_VARIANT` if the variant is malformed. An out-of-bound `id` is also considered
   // a malformed variant because it is read from the corresponding variant value.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala

@@ -1933,13 +1933,24 @@ object SchemaOfVariant {
         val field = v.getFieldAtIndex(i)
         fields(i) = StructField(field.key, schemaOf(field.value))
       }
-      // According to the variant spec, object fields must be sorted alphabetically. So we don't
-      // have to sort, but just need to validate they are sorted.
+      var utf8Sorted = true
+      var utf16Sorted = true
+      var previousKey = if (size > 0) fields(0).name else null
+      var previousKeyBytes = if (size > 0) VariantUtil.encodeKey(previousKey) else null
       for (i <- 1 until size) {
-        if (fields(i - 1).name >= fields(i).name) {
-          throw new SparkRuntimeException("MALFORMED_VARIANT", Map.empty)
-        }
+        val currentKey = fields(i).name
+        val currentKeyBytes = VariantUtil.encodeKey(currentKey)
+        utf8Sorted &&= VariantUtil.compareKeys(previousKeyBytes, currentKeyBytes) < 0
+        utf16Sorted &&= previousKey.compareTo(currentKey) < 0
+        previousKey = currentKey
+        previousKeyBytes = currentKeyBytes
+      }
+      if (!utf8Sorted && !utf16Sorted) {
+        throw new SparkRuntimeException("MALFORMED_VARIANT", Map.empty)
       }
+      // `mergeSchema` expects StructType fields in Java String order. Older Spark values already
+      // use that order, while spec-compliant values need to be reordered after validation.
+      java.util.Arrays.sort(fields, JsonInferSchema.structFieldComparator)
       StructType(fields)
     case Type.ARRAY =>
       var elementType: DataType = NullType

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala

@@ -548,6 +548,57 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper {
     testVariantGet(json, "$." + numKeys, IntegerType, null)
   }
 
+  test("SPARK-58949: object keys use unsigned UTF-8 order") {
+    val bmpKey = new String(Character.toChars(65535))
+    val supplementaryKey = new String(Character.toChars(0x10000))
+    val quote = 34.toChar.toString
+    val asciiFields = (0 until 32).map(i => quote + i + quote + ":" + i)
+    val objectJson = (asciiFields ++ Seq(
+      quote + supplementaryKey + quote + ":99",
+      quote + bmpKey + quote + ":98")).mkString("{", ",", "}")
+
+    val variant = VariantBuilder.parseJson(objectJson, false)
+    assert(variant.getFieldAtIndex(32).key === bmpKey)
+    assert(variant.getFieldAtIndex(33).key === supplementaryKey)
+    assert(variant.getFieldByKey(bmpKey).getLong === 98L)
+    assert(variant.getFieldByKey(supplementaryKey).getLong === 99L)
+    assert(variant.getFieldByKey("missing") === null)
+
+    val nestedJson = "{" + quote + "nested" + quote + ":" + objectJson + "}"
+    val nested = VariantBuilder.parseJson(nestedJson, false)
+      .getFieldByKey("nested")
+    assert(nested.getFieldAtIndex(32).key === bmpKey)
+    assert(nested.getFieldByKey(supplementaryKey).getLong === 99L)
+
+    // Reorder the last two field entries to reproduce the UTF-16 order written by older Spark.
+    val legacyValue = variant.getValue.clone()
+    handleObject[Unit](legacyValue, 0,
+      (size, idSize, offsetSize, idStart, offsetStart, _dataStart) => {
+        def swap(start: Int, width: Int): Unit = {
+          val left = start + (size - 2) * width
+          val right = left + width
+          val leftValue = readUnsigned(legacyValue, left, width)
+          val rightValue = readUnsigned(legacyValue, right, width)
+          writeLong(legacyValue, left, rightValue, width)
+          writeLong(legacyValue, right, leftValue, width)
+        }
+        swap(idStart, idSize)
+        swap(offsetStart, offsetSize)
+      })
+    val legacy = new Variant(legacyValue, variant.getMetadata)
+    assert(legacy.getFieldAtIndex(32).key === supplementaryKey)
+    assert(legacy.getFieldByKey("31").getLong === 31L)
+    assert(legacy.getFieldByKey(bmpKey).getLong === 98L)
+    assert(legacy.getFieldByKey("missing") === null)
+
+    val expectedSchemaNames = ((0 until 32).map(_.toString).sorted ++
+      Seq(supplementaryKey, bmpKey)).toArray
+    Seq(variant, legacy).foreach { v =>
+      val schema = SchemaOfVariant.schemaOf(v).asInstanceOf[StructType]
+      assert(schema.fieldNames === expectedSchemaNames)
+    }
+  }
+
   test("variant_get timestamp") {
     DateTimeTestUtils.outstandingZoneIds.foreach { zid =>
       withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> zid.getId) {

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

@@ -365,16 +365,23 @@ class InferVariantShreddingSchema(val schema: StructType) {
     v.getType match {
       case Type.OBJECT =>
         val size = v.objectSize()
-        // Validate fields are sorted (per variant spec)
+        var utf8Sorted = true
+        var utf16Sorted = true
+        var previousKey = if (size > 0) v.getFieldAtIndex(0).key else null
+        var previousKeyBytes = if (size > 0) VariantUtil.encodeKey(previousKey) else null
         for (i <- 1 until size) {
-          val prevKey = v.getFieldAtIndex(i - 1).key
-          val currKey = v.getFieldAtIndex(i).key
-          if (prevKey >= currKey) {
-            throw new SparkRuntimeException(
-              errorClass = "MALFORMED_VARIANT",
-              messageParameters = Map.empty
-            )
-          }
+          val currentKey = v.getFieldAtIndex(i).key
+          val currentKeyBytes = VariantUtil.encodeKey(currentKey)
+          utf8Sorted &&= VariantUtil.compareKeys(previousKeyBytes, currentKeyBytes) < 0
+          utf16Sorted &&= previousKey.compareTo(currentKey) < 0
+          previousKey = currentKey
+          previousKeyBytes = currentKeyBytes
+        }
+        if (!utf8Sorted && !utf16Sorted) {
+          throw new SparkRuntimeException(
+            errorClass = "MALFORMED_VARIANT",
+            messageParameters = Map.empty
+          )
         }
 
         // Process each field

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

@@ -764,6 +764,45 @@ class VariantInferShreddingSuite extends SharedSparkSession with ParquetTest {
     checkAnswer(spark.read.parquet(dir.getAbsolutePath), df.collect())
   }
 
+  testWithTempDir(
+      "SPARK-58949: infer shredding schema from canonical and legacy fields") { dir =>
+    val bmpKey = new String(Character.toChars(65535))
+    val supplementaryKey = new String(Character.toChars(0x10000))
+    val quote = 34.toChar.toString
+    val json = "{" + quote + supplementaryKey + quote + ":1," +
+      quote + bmpKey + quote + ":2}"
+    val canonical = VariantBuilder.parseJson(json, false)
+    val legacyValue = canonical.getValue.clone()
+    VariantUtil.handleObject[Unit](legacyValue, 0,
+      (size, idSize, offsetSize, idStart, offsetStart, _dataStart) => {
+        def swap(start: Int, width: Int): Unit = {
+          val left = start + (size - 2) * width
+          val right = left + width
+          val leftValue = VariantUtil.readUnsigned(legacyValue, left, width)
+          val rightValue = VariantUtil.readUnsigned(legacyValue, right, width)
+          VariantUtil.writeLong(legacyValue, left, rightValue, width)
+          VariantUtil.writeLong(legacyValue, right, leftValue, width)
+        }
+        swap(idStart, idSize)
+        swap(offsetStart, offsetSize)
+      })
+    val canonicalValue = canonical.getValue
+    val metadata = canonical.getMetadata
+    val rdd = spark.sparkContext.parallelize(0 until 20, 1).map { i =>
+      val value = if (i % 2 == 0) canonicalValue else legacyValue
+      InternalRow(new VariantVal(value, metadata))
+    }
+    val writeSchema = DataType.fromDDL("struct<v variant>").asInstanceOf[StructType]
+    val df = Dataset.ofRows(spark, LogicalRDD(DataTypeUtils.toAttributes(writeSchema), rdd)(spark))
+
+    df.write.mode("overwrite").parquet(dir.getAbsolutePath)
+    val expected = StructType(Seq(
+      StructField(supplementaryKey, LongType),
+      StructField(bmpKey, LongType)))
+    checkFileSchema(expected, dir)
+    assert(spark.read.parquet(dir.getAbsolutePath).count() === 20)
+  }
+
   testWithTempDir("special characters in field names - dots") { dir =>
     val df = spark.sql(
       """