GH-3735: Use unsigned UTF-8 byte order for Variant object field keys

parquet-variant/src/main/java/org/apache/parquet/variant/Variant.java

@@ -273,24 +273,41 @@ public Variant getFieldByKey(String key) {
         }
       }
     } else {
-      int low = 0;
-      int high = info.numElements - 1;
-      while (low <= high) {
-        // Use unsigned right shift to compute the middle of `low` and `high`. This is not only a
-        // performance optimization, because it can properly handle the case where `low + high`
-        // overflows int.
-        int mid = (low + high) >>> 1;
-        int midId = VariantUtil.readUnsignedLittleEndian(value, idStart + info.idSize * mid, info.idSize);
-        String midKey = getMetadataKeyCached(midId);
-        int cmp = midKey.compareTo(key);
-        if (cmp < 0) {
-          low = mid + 1;
-        } else if (cmp > 0) {
-          high = mid - 1;
-        } else {
-          int offset = VariantUtil.readUnsignedLittleEndian(
-              value, offsetStart + info.offsetSize * mid, info.offsetSize);
-          return childVariant(VariantUtil.slice(value, dataStart + offset));
+      // Encode the lookup key once, outside the loop, rather than on every comparison.
+      byte[] keyBytes = VariantUtil.encodeKey(key);
+      // UTF-8 and UTF-16 orders can only differ for keys containing a code unit at or above
+      // U+D800; for all other keys a single search covers both orders.
+      int numAttempts = 1;
+      for (int i = 0; i < key.length(); ++i) {
+        if (key.charAt(i) >= Character.MIN_SURROGATE) {
+          numAttempts = 2;
+          break;
+        }
+      }
+      // Search in the spec's unsigned UTF-8 byte order first, then retry in the UTF-16 order
+      // written by older versions, so objects written before the ordering fix remain readable.
+      for (int attempt = 0; attempt < numAttempts; ++attempt) {
+        int low = 0;
+        int high = info.numElements - 1;
+        while (low <= high) {
+          // Use unsigned right shift to compute the middle of `low` and `high`. This is not only a
+          // performance optimization, because it can properly handle the case where `low + high`
+          // overflows int.
+          int mid = (low + high) >>> 1;
+          int midId = VariantUtil.readUnsignedLittleEndian(value, idStart + info.idSize * mid, info.idSize);
+          String midKey = getMetadataKeyCached(midId);
+          int cmp = attempt == 0
+              ? VariantUtil.compareKeys(VariantUtil.encodeKey(midKey), keyBytes)
+              : midKey.compareTo(key);
+          if (cmp < 0) {
+            low = mid + 1;
+          } else if (cmp > 0) {
+            high = mid - 1;
+          } else {
+            int offset = VariantUtil.readUnsignedLittleEndian(
+                value, offsetStart + info.offsetSize * mid, info.offsetSize);
+            return childVariant(VariantUtil.slice(value, dataStart + offset));
+          }
         }
       }
     }

parquet-variant/src/main/java/org/apache/parquet/variant/VariantBuilder.java

@@ -679,6 +679,12 @@ static final class FieldEntry implements Comparable<FieldEntry> {
     final int offset;
     int valueSize = 0;
 
+    /**
+     * Lazy cache of the UTF-8 encoding of `key`, which sorting an object compares O(log n) times
+     * per entry. Encoded on demand so single-field objects, which are never compared, skip it.
+     */
+    private byte[] keyBytes;
+
     FieldEntry(String key, int id, int offset) {
       this.key = key;
       this.id = id;
@@ -689,9 +695,16 @@ void updateValueSize(int size) {
       valueSize = size;
     }
 
+    private byte[] keyBytes() {
+      if (keyBytes == null) {
+        keyBytes = VariantUtil.encodeKey(key);
+      }
+      return keyBytes;
+    }
+
     @Override
     public int compareTo(FieldEntry other) {
-      return key.compareTo(other.key);
+      return VariantUtil.compareKeys(keyBytes(), other.keyBytes());
     }
   }
 

parquet-variant/src/main/java/org/apache/parquet/variant/VariantUtil.java

@@ -301,6 +301,31 @@ static int readUnsigned(ByteBuffer bytes, int pos, int numBytes) {
     return result;
   }
 
+  /**
+   * Encodes an object field key to the UTF-8 bytes that {@link #compareKeys} orders. Callers that
+   * compare the same key repeatedly - sorting an object, or binary-searching it for one key -
+   * should encode it once and reuse the result rather than re-encoding per comparison.
+   */
+  static byte[] encodeKey(String key) {
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  /**
+   * Compares two object field keys, given their UTF-8 encodings, by unsigned lexicographic byte
+   * order, as required by the Variant spec for object field ordering.
+   *
+   * <p>This intentionally differs from {@link String#compareTo}, which compares UTF-16 code
+   * units. The two orderings agree for all keys in the Basic Multilingual Plane but diverge for
+   * supplementary-plane characters (U+10000 and above): {@code String#compareTo} orders a leading
+   * high surrogate (0xD800-0xDBFF) before code points in U+E000..U+FFFF, whereas UTF-8 byte order
+   * (and the spec) orders them after. Using UTF-16 order here would produce objects whose field
+   * ids are mis-sorted relative to the spec, breaking binary-search lookups by any reader that
+   * follows the spec's UTF-8 byte ordering.
+   */
+  static int compareKeys(byte[] a, byte[] b) {
+    return Arrays.compareUnsigned(a, b);
+  }
+
   /**
    * Fast little-endian unsigned read using bulk ByteBuffer operations.
    * Requires the buffer to have {@link java.nio.ByteOrder#LITTLE_ENDIAN} byte order.

parquet-variant/src/test/java/org/apache/parquet/variant/TestVariantObjectBuilder.java

@@ -85,6 +85,131 @@ public void testLargeObjectBuilder() {
     });
   }
 
+  /**
+   * Object field keys must be ordered by the unsigned byte order of their UTF-8 encoding, not by
+   * {@link String#compareTo} (UTF-16 code-unit order). The two orderings disagree for
+   * supplementary-plane keys: U+FFFF encodes to UTF-8 {@code EF BF BF} and U+10000 to
+   * {@code F0 90 80 80}, so U+FFFF must sort first; but in UTF-16 the leading high surrogate
+   * 0xD800 of U+10000 sorts before 0xFFFF, which would wrongly put U+10000 first. See
+   * {@link VariantUtil#compareKeys}.
+   */
+  @Test
+  public void testObjectKeysSortedByUtf8ByteOrder() {
+    String bmpKey = "￿"; // U+FFFF -> UTF-8 EF BF BF
+    String supplementaryKey = new String(Character.toChars(0x10000)); // UTF-8 F0 90 80 80
+
+    VariantBuilder b = new VariantBuilder();
+    VariantObjectBuilder o = b.startObject();
+    // Appended in the "wrong" order on purpose, to prove the builder sorts rather than
+    // preserving insertion order.
+    o.appendKey(supplementaryKey);
+    o.appendLong(2);
+    o.appendKey(bmpKey);
+    o.appendLong(1);
+    b.endObject();
+
+    VariantTestUtil.testVariant(b.build(), v -> {
+      VariantTestUtil.checkType(v, VariantUtil.OBJECT, Variant.Type.OBJECT);
+      assertThat(v.numObjectElements()).isEqualTo(2);
+      // UTF-8 byte order: EF BF BF < F0 90 80 80, so the BMP key comes first.
+      assertThat(v.getFieldAtIndex(0).key).isEqualTo(bmpKey);
+      assertThat(v.getFieldAtIndex(1).key).isEqualTo(supplementaryKey);
+      assertThat(v.getFieldByKey(bmpKey).getLong()).isEqualTo(1);
+      assertThat(v.getFieldByKey(supplementaryKey).getLong()).isEqualTo(2);
+    });
+  }
+
+  /**
+   * A large object (>= BINARY_SEARCH_THRESHOLD) that mixes ASCII keys with U+FFFF and a
+   * supplementary-plane key, exercising the reader's binary-search path in
+   * {@link Variant#getFieldByKey}. The binary search must use the same UTF-8 byte ordering as the
+   * builder's sort; with a UTF-16 comparator on the read side, the supplementary key would be
+   * mis-navigated and not found.
+   */
+  @Test
+  public void testLargeObjectBinarySearchWithSupplementaryKey() {
+    String bmpKey = "￿"; // UTF-8 EF BF BF
+    String supplementaryKey = new String(Character.toChars(0x10000)); // UTF-8 F0 90 80 80
+
+    VariantBuilder b = new VariantBuilder();
+    VariantObjectBuilder o = b.startObject();
+    for (int i = 0; i < 40; i++) { // well above BINARY_SEARCH_THRESHOLD (32)
+      o.appendKey(String.format("a%03d", i));
+      o.appendLong(i);
+    }
+    o.appendKey(bmpKey);
+    o.appendLong(998);
+    o.appendKey(supplementaryKey);
+    o.appendLong(999);
+    b.endObject();
+
+    VariantTestUtil.testVariant(b.build(), v -> {
+      assertThat(v.numObjectElements()).isEqualTo(42);
+      assertThat(v.getFieldByKey(bmpKey)).isNotNull();
+      assertThat(v.getFieldByKey(bmpKey).getLong()).isEqualTo(998);
+      assertThat(v.getFieldByKey(supplementaryKey)).isNotNull();
+      assertThat(v.getFieldByKey(supplementaryKey).getLong()).isEqualTo(999);
+      assertThat(v.getFieldByKey("a037").getLong()).isEqualTo(37);
+    });
+  }
+
+  /**
+   * Objects written before the ordering fix sorted field ids by {@link String#compareTo} (UTF-16
+   * order). {@link Variant#getFieldByKey} must still find keys in such objects: when a key
+   * contains a code unit at or above U+D800, the lookup retries the binary search in UTF-16 order
+   * after the spec's UTF-8 order fails.
+   */
+  @Test
+  public void testLegacyUtf16OrderedObjectLookup() {
+    String bmpKey = "￿"; // UTF-8 EF BF BF
+    String supplementaryKey = new String(Character.toChars(0x10000)); // UTF-8 F0 90 80 80
+
+    VariantBuilder b = new VariantBuilder();
+    VariantObjectBuilder o = b.startObject();
+    for (int i = 0; i < 40; i++) {
+      o.appendKey(String.format("a%03d", i));
+      o.appendLong(i);
+    }
+    o.appendKey(bmpKey);
+    o.appendLong(998);
+    o.appendKey(supplementaryKey);
+    o.appendLong(999);
+    b.endObject();
+    Variant canonical = b.build();
+
+    // Reproduce the layout written by older versions: swap the id and offset entries of the last
+    // two fields, so the supplementary key precedes the BMP key (UTF-16 order).
+    ByteBuffer valueBuffer = canonical.getValueBuffer().duplicate();
+    byte[] legacyValue = new byte[valueBuffer.remaining()];
+    valueBuffer.get(legacyValue);
+    VariantUtil.ObjectInfo info =
+        VariantUtil.getObjectInfo(ByteBuffer.wrap(legacyValue).order(ByteOrder.LITTLE_ENDIAN));
+    swapLastTwoEntries(legacyValue, info.idStartOffset, info.idSize, info.numElements);
+    swapLastTwoEntries(legacyValue, info.offsetStartOffset, info.offsetSize, info.numElements);
+    Variant legacy = new Variant(ByteBuffer.wrap(legacyValue), canonical.getMetadataBuffer());
+
+    assertThat(legacy.getFieldAtIndex(40).key).isEqualTo(supplementaryKey);
+    assertThat(legacy.getFieldAtIndex(41).key).isEqualTo(bmpKey);
+    // ASCII keys are found by the first (UTF-8 order) search.
+    assertThat(legacy.getFieldByKey("a037").getLong()).isEqualTo(37);
+    // Keys at or above U+D800 are found by the UTF-16 order retry.
+    assertThat(legacy.getFieldByKey(bmpKey).getLong()).isEqualTo(998);
+    assertThat(legacy.getFieldByKey(supplementaryKey).getLong()).isEqualTo(999);
+    // Absent keys stay absent after both attempts.
+    assertThat(legacy.getFieldByKey("missing")).isNull();
+    assertThat(legacy.getFieldByKey(new String(Character.toChars(0x10001)))).isNull();
+  }
+
+  private static void swapLastTwoEntries(byte[] bytes, int start, int width, int numElements) {
+    int left = start + (numElements - 2) * width;
+    int right = left + width;
+    ByteBuffer buffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN);
+    int leftValue = VariantUtil.readUnsignedLittleEndian(buffer, left, width);
+    int rightValue = VariantUtil.readUnsignedLittleEndian(buffer, right, width);
+    VariantUtil.writeLong(bytes, left, rightValue, width);
+    VariantUtil.writeLong(bytes, right, leftValue, width);
+  }
+
   @Test
   public void testMixedObjectBuilder() {
     VariantBuilder b = new VariantBuilder();