SavicStefan commented on code in PR #59669:
URL: https://github.com/apache/spark/pull/59669#discussion_r4452761200
##########
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) {
Review Comment:
yeah, makes 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]
