diff --git a/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariant.java b/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariant.java index 254a742ededbeb..90e61c159df304 100644 --- a/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariant.java +++ b/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariant.java @@ -37,10 +37,12 @@ import static org.apache.flink.types.variant.BinaryVariantUtil.VERSION; import static org.apache.flink.types.variant.BinaryVariantUtil.VERSION_MASK; import static org.apache.flink.types.variant.BinaryVariantUtil.checkIndex; +import static org.apache.flink.types.variant.BinaryVariantUtil.encodeKey; import static org.apache.flink.types.variant.BinaryVariantUtil.getMetadataKey; import static org.apache.flink.types.variant.BinaryVariantUtil.handleArray; import static org.apache.flink.types.variant.BinaryVariantUtil.handleObject; import static org.apache.flink.types.variant.BinaryVariantUtil.malformedVariant; +import static org.apache.flink.types.variant.BinaryVariantUtil.metadataKeyEquals; import static org.apache.flink.types.variant.BinaryVariantUtil.microsToInstant; import static org.apache.flink.types.variant.BinaryVariantUtil.nanosToInstant; import static org.apache.flink.types.variant.BinaryVariantUtil.readUnsigned; @@ -344,9 +346,10 @@ private BinaryVariant getFieldByKey(String key) { // Use linear search for a short list. Switch to binary search when the length // reaches `BINARY_SEARCH_THRESHOLD`. if (size < BINARY_SEARCH_THRESHOLD) { + byte[] keyBytes = encodeKey(key); for (int i = 0; i < size; ++i) { int id = readUnsigned(value, idStart + idSize * i, idSize); - if (key.equals(getMetadataKey(metadata, id))) { + if (metadataKeyEquals(metadata, id, keyBytes)) { int offset = readUnsigned( value, offsetStart + offsetSize * i, offsetSize); @@ -362,6 +365,8 @@ private BinaryVariant getFieldByKey(String key) { // handle the case where `low + high` overflows int. int mid = (low + high) >>> 1; int id = readUnsigned(value, idStart + idSize * mid, idSize); + // Compare decoded Strings here. The builder sorts fields by + // String.compareTo, which is not UTF-8 byte order. int cmp = getMetadataKey(metadata, id).compareTo(key); if (cmp < 0) { low = mid + 1; diff --git a/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantInternalBuilder.java b/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantInternalBuilder.java index 3ed2e634412152..4ac26cb2b46440 100644 --- a/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantInternalBuilder.java +++ b/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantInternalBuilder.java @@ -74,6 +74,7 @@ import static org.apache.flink.types.variant.BinaryVariantUtil.VERSION; import static org.apache.flink.types.variant.BinaryVariantUtil.arrayHeader; import static org.apache.flink.types.variant.BinaryVariantUtil.checkIndex; +import static org.apache.flink.types.variant.BinaryVariantUtil.encodeKey; import static org.apache.flink.types.variant.BinaryVariantUtil.fitsVariantDecimal; import static org.apache.flink.types.variant.BinaryVariantUtil.getMetadataKey; import static org.apache.flink.types.variant.BinaryVariantUtil.handleArray; @@ -477,7 +478,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; } diff --git a/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantUtil.java b/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantUtil.java index 546b5d4bcdde32..c0812989fa6336 100644 --- a/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantUtil.java +++ b/flink-core/src/main/java/org/apache/flink/types/variant/BinaryVariantUtil.java @@ -850,6 +850,31 @@ public static T handleArray(byte[] value, int pos, ArrayHandler handler) * considered a malformed variant because it is read from the corresponding variant value. */ public static String getMetadataKey(byte[] metadata, int id) { + long range = metadataKeyRange(metadata, id); + return new String(metadata, keyStart(range), keyLength(range), StandardCharsets.UTF_8); + } + + /** Encodes an object key the way the variant metadata stores it. */ + static byte[] encodeKey(String key) { + return key.getBytes(StandardCharsets.UTF_8); + } + + /** Whether the key at {@code id} in the variant metadata is the UTF-8 encoded {@code key}. */ + static boolean metadataKeyEquals(byte[] metadata, int id, byte[] key) { + long range = metadataKeyRange(metadata, id); + int start = keyStart(range); + int length = keyLength(range); + return length == key.length + && Arrays.equals(metadata, start, start + length, key, 0, length); + } + + /** + * Locates the key at {@code id} in the variant metadata. The range is packed into a long so + * that a lookup allocates nothing. + * + * @return the start index of the key in the high 32 bits and its length in the low 32 bits + */ + private static long metadataKeyRange(byte[] metadata, int id) { checkIndex(0, metadata.length); // Extracts the highest 2 bits in the metadata header to determine the integer size of the // offset list. @@ -867,8 +892,15 @@ public static String getMetadataKey(byte[] metadata, int id) { throw malformedVariant(); } checkIndex(stringStart + nextOffset - 1, metadata.length); - return new String( - metadata, stringStart + offset, nextOffset - offset, StandardCharsets.UTF_8); + return ((long) (stringStart + offset) << 32) | (nextOffset - offset); + } + + private static int keyStart(long range) { + return (int) (range >>> 32); + } + + private static int keyLength(long range) { + return (int) range; } /** diff --git a/flink-core/src/test/java/org/apache/flink/types/variant/BinaryVariantTest.java b/flink-core/src/test/java/org/apache/flink/types/variant/BinaryVariantTest.java index e3725c1d202f55..4dd84d1d0db9df 100644 --- a/flink-core/src/test/java/org/apache/flink/types/variant/BinaryVariantTest.java +++ b/flink-core/src/test/java/org/apache/flink/types/variant/BinaryVariantTest.java @@ -488,6 +488,40 @@ void testNonAsciiStringsAndFieldNames() { + "\"}}"); } + @Test + void testGetFieldWithUnpairedSurrogate() { + // UTF-8 has no unpaired surrogates. The builder writes one as '?', and a lookup encodes its + // key the same way. + final Variant object = builder.object().add("a\uD800", builder.of(1)).build(); + + assertThat(object.getFieldNames()).containsExactly("a?"); + assertThat(object.getField("a\uD800").getInt()).isEqualTo(1); + assertThat(object.getField("a?").getInt()).isEqualTo(1); + } + + @Test + void testGetFieldWithEmptyKey() { + final Variant object = + builder.object().add("", builder.of(1)).add("b", builder.of(2)).build(); + + assertThat(object.getField("").getInt()).isEqualTo(1); + assertThat(object.getField("b").getInt()).isEqualTo(2); + } + + @Test + void testGetFieldWithMalformedUtf8Key() { + final BinaryVariant object = + (BinaryVariant) builder.object().add("a", builder.of(1)).build(); + // The key is the last metadata byte. 0xFF never occurs in UTF-8, so no lookup key encodes + // to it, even though getFieldNames decodes it to U+FFFD. + final byte[] metadata = object.getMetadata().clone(); + metadata[metadata.length - 1] = (byte) 0xFF; + final Variant malformed = new BinaryVariant(object.getValue(), metadata); + + assertThat(malformed.getFieldNames()).containsExactly("�"); + assertThat(malformed.getField("�")).isNull(); + } + @Test void testVariantException() { assertThatThrownBy(() -> new BinaryVariant(new byte[0], new byte[0]))