SavicStefan commented on code in PR #59669:
URL: https://github.com/apache/spark/pull/59669#discussion_r4389463893
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -555,18 +654,254 @@ private void appendVariantImpl(byte[] value, byte[]
metadata, int pos) {
int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
int elementPos = dataStart + offset;
offsets.add(writePos - start);
- appendVariantImpl(value, metadata, elementPos);
+ appendVariantImpl(value, metadata, elementPos,
needCanonicalization);
}
finishWritingArray(start, offsets);
return null;
});
break;
+ default:
+ if (needCanonicalization) {
+ appendCanonicalizedScalar(value, pos);
+ } else {
+ shallowAppendVariantImpl(value, pos);
+ }
+ break;
+ }
+ }
+
+ // Append the scalar at `value[pos...]` in canonical form. Keep in sync with
the scalar cases
+ // of `isValueCanonical`.
+ private void appendCanonicalizedScalar(byte[] value, int pos) {
+ switch (getType(value, pos)) {
+ case LONG:
+ appendLong(getLong(value, pos));
+ break;
+ case DECIMAL: {
+ BigDecimal stripped = getDecimal(value, pos);
+ if (decimalPromotesToLong(stripped)) {
+ appendLong(stripped.longValue());
+ } else {
+ appendDecimal(canonicalDecimalForm(stripped));
+ }
+ break;
+ }
+ case FLOAT:
+ appendFloat(canonicalizeFloat(getFloat(value, pos)));
+ break;
+ case DOUBLE:
+ appendDouble(canonicalizeDouble(getDouble(value, pos)));
+ break;
+ case STRING:
+ if (isCanonicalString(value, pos)) {
+ shallowAppendVariantImpl(value, pos);
+ } else {
+ appendString(getString(value, pos));
+ }
+ break;
default:
shallowAppendVariantImpl(value, pos);
break;
}
}
+ private void buildCanonicalized(byte[] value, byte[] metadata, int pos) {
+ int offsetSize = getMetadataOffsetSize(metadata);
+ int numKeys = getMetadataNumKeys(metadata, offsetSize);
+ getMetadataStringStart(metadata, numKeys, offsetSize);
+ boolean[] seen = new boolean[numKeys];
+ ArrayList<byte[]> keys = new ArrayList<>();
+ collectAllObjectKeys(value, metadata, pos, seen, keys);
+ // Sort by UTF-8 bytes; String order differs for characters above U+FFFF.
+ keys.sort(VariantUtil::compareKeys);
+ for (byte[] key : keys) {
+ addKey(new String(key, StandardCharsets.UTF_8));
+ }
+ appendVariantImpl(value, metadata, pos, true);
+ }
+
+ private static void collectAllObjectKeys(
+ byte[] value, byte[] metadata, int pos, boolean[] seen,
ArrayList<byte[]> keys) {
+ checkIndex(pos, value.length);
+ switch (value[pos] & BASIC_TYPE_MASK) {
+ case OBJECT:
+ handleObject(value, pos, (size, idSize, offsetSize, idStart,
offsetStart, dataStart) -> {
+ for (int fieldIndex = 0; fieldIndex < size; ++fieldIndex) {
+ int keyId = readUnsigned(value, idStart + idSize * fieldIndex,
idSize);
+ int offset = readUnsigned(value, offsetStart + offsetSize *
fieldIndex, offsetSize);
+ if (keyId >= seen.length) {
+ throw malformedVariant();
+ }
+ if (!seen[keyId]) {
+ seen[keyId] = true;
+ keys.add(encodeKey(getMetadataKey(metadata, keyId)));
+ }
+ collectAllObjectKeys(value, metadata, dataStart + offset, seen,
keys);
+ }
+ return null;
+ });
+ break;
+ case ARRAY:
+ handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> {
+ for (int elementIndex = 0; elementIndex < size; ++elementIndex) {
+ int offset = readUnsigned(value, offsetStart + offsetSize *
elementIndex, offsetSize);
+ collectAllObjectKeys(value, metadata, dataStart + offset, seen,
keys);
+ }
+ return null;
+ });
+ break;
+ default:
+ break;
+ }
+ }
+
+ // Check canonical value encoding and mark key IDs in referenced as they are
visited.
+ // On false, referenced may be incomplete. Keep these checks in sync with
+ // appendCanonicalizedScalar, canonical object/array rebuilding in
appendVariantImpl,
+ // finishWritingObject, and finishWritingArray.
+ private static boolean isValueCanonical(byte[] value, int pos, boolean[]
referenced) {
+ int encodedSize = valueSize(value, pos);
+ if (encodedSize <= 0 || (long) pos + encodedSize > value.length) {
+ return false;
+ }
+ switch (getType(value, pos)) {
+ case OBJECT:
+ return handleObject(value, pos,
+ (size, idSize, offsetSize, idStart, offsetStart, dataStart) -> {
+ int previousId = -1;
+ int dataSize = 0;
+ for (int fieldIndex = 0; fieldIndex < size; ++fieldIndex) {
+ int keyId = readUnsigned(value, idStart + idSize * fieldIndex,
idSize);
+ if (keyId <= previousId || keyId >= referenced.length) {
+ return false;
+ }
+ previousId = keyId;
+ referenced[keyId] = true;
+ int offset = readUnsigned(value, offsetStart + offsetSize *
fieldIndex, offsetSize);
+ if (offset != dataSize) {
+ return false;
+ }
+ int elementPos = dataStart + offset;
+ if (!isValueCanonical(value, elementPos, referenced)) {
+ return false;
+ }
+ dataSize += valueSize(value, elementPos);
+ }
+ int lastOffset = readUnsigned(value, offsetStart + offsetSize *
size, offsetSize);
+ if (lastOffset != dataSize) {
+ return false;
+ }
+ int maxId = size == 0 ? 0 : previousId;
+ return value[pos] ==
+ objectHeader(size > U8_MAX, getIntegerSize(maxId),
getIntegerSize(dataSize));
+ });
+ case ARRAY:
+ return handleArray(value, pos, (size, offsetSize, offsetStart,
dataStart) -> {
+ int dataSize = 0;
+ for (int elementIndex = 0; elementIndex < size; ++elementIndex) {
+ int offset = readUnsigned(value, offsetStart + offsetSize *
elementIndex, offsetSize);
+ if (offset != dataSize) {
+ return false;
+ }
+ int elementPos = dataStart + offset;
+ if (!isValueCanonical(value, elementPos, referenced)) {
+ return false;
+ }
+ dataSize += valueSize(value, elementPos);
+ }
+ int lastOffset = readUnsigned(value, offsetStart + offsetSize *
size, offsetSize);
+ return lastOffset == dataSize &&
+ value[pos] == arrayHeader(size > U8_MAX,
getIntegerSize(dataSize));
+ });
+ case LONG:
+ return getTypeInfo(value, pos) == minimalLongTypeInfo(getLong(value,
pos));
+ case DECIMAL: {
+ BigDecimal original = getDecimalWithOriginalScale(value, pos);
+ BigDecimal stripped = original.stripTrailingZeros();
+ if (decimalPromotesToLong(stripped)) {
+ return false;
+ }
+ BigDecimal canonical = canonicalDecimalForm(stripped);
+ return original.scale() == canonical.scale() &&
+ getTypeInfo(value, pos) == minimalDecimalTypeInfo(canonical);
+ }
+ case FLOAT: {
+ int rawBits = (int) readLong(value, pos + 1, 4);
+ float number = Float.intBitsToFloat(rawBits);
+ return rawBits == Float.floatToIntBits(canonicalizeFloat(number));
+ }
+ case DOUBLE: {
+ long rawBits = readLong(value, pos + 1, 8);
+ double number = Double.longBitsToDouble(rawBits);
+ return rawBits == Double.doubleToLongBits(canonicalizeDouble(number));
+ }
+ case STRING:
+ return isCanonicalString(value, pos);
+ default:
+ // Other currently supported types are copied unchanged by the rebuild.
+ // New types with alternative encodings need an explicit canonicality
check.
+ return true;
+ }
+ }
+
+ // Return whether the string at `value[pos...]` is valid UTF-8 and uses
short-string
+ // encoding whenever it fits.
+ private static boolean isCanonicalString(byte[] value, int pos) {
+ int encodedSize = valueSize(value, pos);
+ if (encodedSize <= 0 || (long) pos + encodedSize > value.length) {
+ return false;
+ }
+ boolean shortString = (value[pos] & BASIC_TYPE_MASK) == SHORT_STR;
+ int headerSize = shortString ? 1 : 1 + U32_SIZE;
+ int length = encodedSize - headerSize;
+ return (shortString || length > MAX_SHORT_STR_SIZE) &&
+ isValidUtf8(value, pos + headerSize, length);
+ }
+
+ private static boolean isValidUtf8(byte[] bytes, int start, int length) {
Review Comment:
I changed so we preserve the original bytes, it makes more sense.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
