diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java b/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java index 8996caeb364cf..cefb957f9b976 100644 --- a/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java +++ b/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) { diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java index d9eed06bf3a68..ef00c8379b89c 100644 --- a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java +++ b/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 { 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()); } } diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java b/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java index 681c4038a9a98..184a7b29a9781 100644 --- a/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java +++ b/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. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala index fddb88d953953..c58634cfd8d47 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala +++ b/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 diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala index c69ad520a2418..ed790e643adc0 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala +++ b/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) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/InferVariantShreddingSchema.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/InferVariantShreddingSchema.scala index bbfbbfde0ba4b..98385c796d850 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/InferVariantShreddingSchema.scala +++ b/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 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantInferShreddingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantInferShreddingSuite.scala index cbd3d89c36586..588d0352627b8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantInferShreddingSuite.scala +++ b/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").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( """