SavicStefan commented on code in PR #59669:
URL: https://github.com/apache/spark/pull/59669#discussion_r4452890921
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -555,18 +662,261 @@ 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, canonicalIds);
}
finishWritingArray(start, offsets);
return null;
});
break;
+ default:
+ if (canonicalIds != null) {
+ 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: {
+ int size = valueSize(value, pos);
+ if (size <= 0 || (long) pos + size > value.length) {
+ throw malformedVariant();
+ }
+ int headerSize = (value[pos] & BASIC_TYPE_MASK) == SHORT_STR ? 1 : 1 +
U32_SIZE;
+ appendString(value, pos + headerSize, size - headerSize);
+ break;
+ }
default:
shallowAppendVariantImpl(value, pos);
break;
}
}
+ private void buildCanonicalized(byte[] value, byte[] metadata, int pos) {
+ int offsetSize = getMetadataOffsetSize(metadata);
+ int numKeys = getMetadataNumKeys(metadata, offsetSize);
+ // Validate metadata table bounds before allocating the key-ID mapping.
+ getMetadataStringStart(metadata, numKeys, offsetSize);
+ int[] canonicalIds = new int[numKeys];
+ Arrays.fill(canonicalIds, -1);
+ TreeMap<byte[], Integer> keys = new TreeMap<>(VariantUtil::compareKeys);
+ collectAllObjectKeys(value, metadata, pos, canonicalIds, keys);
+ // Remap the IDs assigned during collection to unsigned byte order without
decoding keys.
+ int[] sortedIds = new int[keys.size()];
+ for (java.util.Map.Entry<byte[], Integer> key : keys.entrySet()) {
+ sortedIds[key.getValue()] = dictionaryKeys.size();
+ dictionaryKeys.add(key.getKey());
+ }
+ for (int i = 0; i < canonicalIds.length; ++i) {
+ if (canonicalIds[i] >= 0) {
+ canonicalIds[i] = sortedIds[canonicalIds[i]];
+ }
+ }
+ appendVariantImpl(value, metadata, pos, canonicalIds);
+ }
+
+ private static void collectAllObjectKeys(
+ byte[] value, byte[] metadata, int pos, int[] canonicalIds,
TreeMap<byte[], Integer> 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 >= canonicalIds.length) {
+ throw malformedVariant();
+ }
+ if (canonicalIds[keyId] < 0) {
+ byte[] key = getMetadataKeyBytes(metadata, keyId);
+ Integer id = keys.get(key);
+ if (id == null) {
+ id = keys.size();
+ keys.put(key, id);
+ }
+ canonicalIds[keyId] = id;
+ }
+ collectAllObjectKeys(value, metadata, dataStart + offset,
canonicalIds, 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,
canonicalIds, 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...]` uses short-string encoding
whenever it fits.
+ private static boolean isCanonicalString(byte[] value, int pos) {
Review Comment:
Yeah, inlined it
--
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]
