harshmotw-db commented on code in PR #59669:
URL: https://github.com/apache/spark/pull/59669#discussion_r4419444748
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -205,6 +205,79 @@ public static Variant pickAtPaths(Variant v, PickNode
tree) {
return builder.result();
}
+ // Return v in canonical form: only referenced dictionary keys, unique and
sorted by unsigned
+ // UTF-8 bytes; object fields stored in key order; minimal integer, decimal,
and offset widths;
+ // strings of at most 63 UTF-8 bytes use short-string encoding; decimal
trailing zeros are
+ // stripped, and integral decimals become longs only when they fit.
Normalize floating-point
+ // -0.0 to 0.0 and NaNs within each type. Preserve string and key bytes,
including invalid UTF-8.
+ // Return v itself if it is already canonical.
+ public static Variant canonicalize(Variant v) {
+ if (isCanonical(v)) {
Review Comment:
When I was exploring this, I had ideas about having an additional field in
the `Variant` class saying `isCanonical` which doCanonicalize would mark as
`true`. This way, we would make sure that each variant is only get
canonicalized once. Checking `isCanonical` could potentially be expensive if
you plan on having this `canonicalize_variant` expression with every equality
operation. In these cases, rows would have `doCanonicalize` run once and
`isCanonical` run over and over again on the same variant that you already know
is canonicalized.
Another benefit of this would be that `isCanonical` is riskier than
`doCanonicalize`. You need to be far more careful when making sure something is
right because you might miss checking something. Think of cases where the
variant spec is modified in the future. The trivial implementation of
`doCanonicalize` would fail at anything it doesn't understand because it needs
to be able to parse and reconstruct everything. `isCanonical` would need to
intentionally be written in a way that it is tolerant of any changes.
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -205,6 +205,79 @@ public static Variant pickAtPaths(Variant v, PickNode
tree) {
return builder.result();
}
+ // Return v in canonical form: only referenced dictionary keys, unique and
sorted by unsigned
+ // UTF-8 bytes; object fields stored in key order; minimal integer, decimal,
and offset widths;
+ // strings of at most 63 UTF-8 bytes use short-string encoding; decimal
trailing zeros are
+ // stripped, and integral decimals become longs only when they fit.
Normalize floating-point
+ // -0.0 to 0.0 and NaNs within each type. Preserve string and key bytes,
including invalid UTF-8.
+ // Return v itself if it is already canonical.
+ public static Variant canonicalize(Variant v) {
+ if (isCanonical(v)) {
+ return v;
+ }
+ return doCanonicalize(v);
+ }
+
+ // Check whether v is already canonically encoded. Return false at the first
+ // non-canonical encoding. Malformed bytes also return false so that the
rebuild decides whether
+ // they fail, but an unreadable header may still throw. Returning false does
not establish
+ // validity.
+ static boolean isCanonical(Variant v) {
+ if (v.pos != 0) {
+ return false;
+ }
+ byte[] value = v.value;
+ byte[] metadata = v.metadata;
+ int offsetSize = getMetadataOffsetSize(metadata);
+ int numKeys = getMetadataNumKeys(metadata, offsetSize);
+ int stringStart = getMetadataStringStart(metadata, numKeys, offsetSize);
+ int lastOffset = readUnsigned(metadata, stringStart - offsetSize,
offsetSize);
+ if (lastOffset > metadata.length - stringStart) {
+ return false;
+ }
+ int canonicalOffsetSize = getIntegerSize(Math.max(lastOffset, numKeys));
+ if (metadata[0] != metadataHeader(canonicalOffsetSize)) {
+ return false;
+ }
+ int previousOffset = readUnsigned(metadata, 1 + offsetSize, offsetSize);
+ if (previousOffset != 0) {
Review Comment:
Should this be called `firstOffset` for readability? You can create another
variable `previousOffset` below for the loop.
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -205,6 +205,79 @@ public static Variant pickAtPaths(Variant v, PickNode
tree) {
return builder.result();
}
+ // Return v in canonical form: only referenced dictionary keys, unique and
sorted by unsigned
+ // UTF-8 bytes; object fields stored in key order; minimal integer, decimal,
and offset widths;
+ // strings of at most 63 UTF-8 bytes use short-string encoding; decimal
trailing zeros are
+ // stripped, and integral decimals become longs only when they fit.
Normalize floating-point
+ // -0.0 to 0.0 and NaNs within each type. Preserve string and key bytes,
including invalid UTF-8.
+ // Return v itself if it is already canonical.
+ public static Variant canonicalize(Variant v) {
+ if (isCanonical(v)) {
+ return v;
+ }
+ return doCanonicalize(v);
+ }
+
+ // Check whether v is already canonically encoded. Return false at the first
+ // non-canonical encoding. Malformed bytes also return false so that the
rebuild decides whether
+ // they fail, but an unreadable header may still throw. Returning false does
not establish
+ // validity.
+ static boolean isCanonical(Variant v) {
+ if (v.pos != 0) {
+ return false;
+ }
+ byte[] value = v.value;
+ byte[] metadata = v.metadata;
+ int offsetSize = getMetadataOffsetSize(metadata);
+ int numKeys = getMetadataNumKeys(metadata, offsetSize);
+ int stringStart = getMetadataStringStart(metadata, numKeys, offsetSize);
+ int lastOffset = readUnsigned(metadata, stringStart - offsetSize,
offsetSize);
+ if (lastOffset > metadata.length - stringStart) {
+ return false;
+ }
+ int canonicalOffsetSize = getIntegerSize(Math.max(lastOffset, numKeys));
+ if (metadata[0] != metadataHeader(canonicalOffsetSize)) {
+ return false;
+ }
+ int previousOffset = readUnsigned(metadata, 1 + offsetSize, offsetSize);
+ if (previousOffset != 0) {
+ return false;
+ }
+ int currentOffset = previousOffset;
+ for (int keyId = 0; keyId < numKeys; ++keyId) {
+ int nextOffset = readUnsigned(metadata, 1 + (keyId + 2) * offsetSize,
offsetSize);
+ if (currentOffset > nextOffset || nextOffset > metadata.length -
stringStart) {
Review Comment:
Isn't the `nextOffset > metadata.length - stringStart` check redundant? We
check `lastOffset > metadata.length - stringStart` and for each offset, we
check that it is greater than the previous offset so the last offset must be
the largest offset (which we have already checked is not too large).
##########
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:
Should we move this function up closer to `isCanonical`?
##########
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) {
+ 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;
+ }
+
+ private static int minimalLongTypeInfo(long value) {
+ if (value == (byte) value) return INT1;
+ if (value == (short) value) return INT2;
+ if (value == (int) value) return INT4;
+ return INT8;
+ }
+
+ private static int minimalDecimalTypeInfo(BigDecimal value) {
+ if (value.scale() <= MAX_DECIMAL4_PRECISION && value.precision() <=
MAX_DECIMAL4_PRECISION) {
+ return DECIMAL4;
+ } else if (value.scale() <= MAX_DECIMAL8_PRECISION &&
+ value.precision() <= MAX_DECIMAL8_PRECISION) {
+ return DECIMAL8;
+ } else {
+ return DECIMAL16;
+ }
+ }
+
+ private static float canonicalizeFloat(float value) {
+ return value == 0.0f ? 0.0f : value;
Review Comment:
I guess you don't check for NaN's because `floatToIntBits` already
canonicalizes NaN's. However, then the function name is misleading because if
someone were to use it, they might expect it to actually canonicalize floats
which it doesn't
##########
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) {
+ 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;
+ }
+
+ private static int minimalLongTypeInfo(long value) {
+ if (value == (byte) value) return INT1;
+ if (value == (short) value) return INT2;
+ if (value == (int) value) return INT4;
+ return INT8;
+ }
+
+ private static int minimalDecimalTypeInfo(BigDecimal value) {
+ if (value.scale() <= MAX_DECIMAL4_PRECISION && value.precision() <=
MAX_DECIMAL4_PRECISION) {
+ return DECIMAL4;
+ } else if (value.scale() <= MAX_DECIMAL8_PRECISION &&
+ value.precision() <= MAX_DECIMAL8_PRECISION) {
+ return DECIMAL8;
+ } else {
+ return DECIMAL16;
+ }
+ }
+
+ private static float canonicalizeFloat(float value) {
+ return value == 0.0f ? 0.0f : value;
+ }
+
+ private static double canonicalizeDouble(double value) {
+ return value == 0.0d ? 0.0d : value;
Review Comment:
ditto
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -205,6 +205,79 @@ public static Variant pickAtPaths(Variant v, PickNode
tree) {
return builder.result();
}
+ // Return v in canonical form: only referenced dictionary keys, unique and
sorted by unsigned
+ // UTF-8 bytes; object fields stored in key order; minimal integer, decimal,
and offset widths;
+ // strings of at most 63 UTF-8 bytes use short-string encoding; decimal
trailing zeros are
+ // stripped, and integral decimals become longs only when they fit.
Normalize floating-point
+ // -0.0 to 0.0 and NaNs within each type. Preserve string and key bytes,
including invalid UTF-8.
+ // Return v itself if it is already canonical.
+ public static Variant canonicalize(Variant v) {
+ if (isCanonical(v)) {
+ return v;
+ }
+ return doCanonicalize(v);
+ }
+
+ // Check whether v is already canonically encoded. Return false at the first
+ // non-canonical encoding. Malformed bytes also return false so that the
rebuild decides whether
+ // they fail, but an unreadable header may still throw. Returning false does
not establish
+ // validity.
+ static boolean isCanonical(Variant v) {
+ if (v.pos != 0) {
+ return false;
+ }
+ byte[] value = v.value;
+ byte[] metadata = v.metadata;
+ int offsetSize = getMetadataOffsetSize(metadata);
+ int numKeys = getMetadataNumKeys(metadata, offsetSize);
+ int stringStart = getMetadataStringStart(metadata, numKeys, offsetSize);
+ int lastOffset = readUnsigned(metadata, stringStart - offsetSize,
offsetSize);
+ if (lastOffset > metadata.length - stringStart) {
Review Comment:
Oh you do have the `metadata.length != stringStart + currentOffset` check
[here](https://github.com/SavicStefan/spark/blob/46215dd2ebf658ec773110a4934732923f4d80b3/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java#L260).
Why don't you just move it here?
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -205,6 +205,79 @@ public static Variant pickAtPaths(Variant v, PickNode
tree) {
return builder.result();
}
+ // Return v in canonical form: only referenced dictionary keys, unique and
sorted by unsigned
+ // UTF-8 bytes; object fields stored in key order; minimal integer, decimal,
and offset widths;
+ // strings of at most 63 UTF-8 bytes use short-string encoding; decimal
trailing zeros are
+ // stripped, and integral decimals become longs only when they fit.
Normalize floating-point
+ // -0.0 to 0.0 and NaNs within each type. Preserve string and key bytes,
including invalid UTF-8.
+ // Return v itself if it is already canonical.
+ public static Variant canonicalize(Variant v) {
+ if (isCanonical(v)) {
+ return v;
+ }
+ return doCanonicalize(v);
+ }
+
+ // Check whether v is already canonically encoded. Return false at the first
+ // non-canonical encoding. Malformed bytes also return false so that the
rebuild decides whether
+ // they fail, but an unreadable header may still throw. Returning false does
not establish
+ // validity.
+ static boolean isCanonical(Variant v) {
+ if (v.pos != 0) {
+ return false;
+ }
+ byte[] value = v.value;
+ byte[] metadata = v.metadata;
+ int offsetSize = getMetadataOffsetSize(metadata);
+ int numKeys = getMetadataNumKeys(metadata, offsetSize);
+ int stringStart = getMetadataStringStart(metadata, numKeys, offsetSize);
+ int lastOffset = readUnsigned(metadata, stringStart - offsetSize,
offsetSize);
+ if (lastOffset > metadata.length - stringStart) {
+ return false;
+ }
+ int canonicalOffsetSize = getIntegerSize(Math.max(lastOffset, numKeys));
Review Comment:
I don't believe this requirement is documented even in the PR description.
If that's the case, it is important to document literally every requirement
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -205,6 +205,79 @@ public static Variant pickAtPaths(Variant v, PickNode
tree) {
return builder.result();
}
+ // Return v in canonical form: only referenced dictionary keys, unique and
sorted by unsigned
+ // UTF-8 bytes; object fields stored in key order; minimal integer, decimal,
and offset widths;
+ // strings of at most 63 UTF-8 bytes use short-string encoding; decimal
trailing zeros are
+ // stripped, and integral decimals become longs only when they fit.
Normalize floating-point
+ // -0.0 to 0.0 and NaNs within each type. Preserve string and key bytes,
including invalid UTF-8.
+ // Return v itself if it is already canonical.
+ public static Variant canonicalize(Variant v) {
+ if (isCanonical(v)) {
+ return v;
+ }
+ return doCanonicalize(v);
+ }
+
+ // Check whether v is already canonically encoded. Return false at the first
+ // non-canonical encoding. Malformed bytes also return false so that the
rebuild decides whether
+ // they fail, but an unreadable header may still throw. Returning false does
not establish
+ // validity.
+ static boolean isCanonical(Variant v) {
Review Comment:
For a function like this, it is very important to document what exactly it
does. The only source for the canonicalization spec that I can find is the PR
description. When a reader reads this signature, how do they know what
canonicalization of variant means? They would need to read the code to
understand its purpose and there is no way to know which parts of the function
are safe to modify (without finding this commit of course).
If you have a documented spec somewhere, can you link it here? If not, can
you describe in detail what the canonical representation of a variant means in
a comment somewhere in this file?
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -205,6 +205,79 @@ public static Variant pickAtPaths(Variant v, PickNode
tree) {
return builder.result();
}
+ // Return v in canonical form: only referenced dictionary keys, unique and
sorted by unsigned
+ // UTF-8 bytes; object fields stored in key order; minimal integer, decimal,
and offset widths;
+ // strings of at most 63 UTF-8 bytes use short-string encoding; decimal
trailing zeros are
+ // stripped, and integral decimals become longs only when they fit.
Normalize floating-point
+ // -0.0 to 0.0 and NaNs within each type. Preserve string and key bytes,
including invalid UTF-8.
+ // Return v itself if it is already canonical.
+ public static Variant canonicalize(Variant v) {
+ if (isCanonical(v)) {
+ return v;
+ }
+ return doCanonicalize(v);
+ }
+
+ // Check whether v is already canonically encoded. Return false at the first
+ // non-canonical encoding. Malformed bytes also return false so that the
rebuild decides whether
+ // they fail, but an unreadable header may still throw. Returning false does
not establish
+ // validity.
+ static boolean isCanonical(Variant v) {
+ if (v.pos != 0) {
+ return false;
+ }
+ byte[] value = v.value;
+ byte[] metadata = v.metadata;
+ int offsetSize = getMetadataOffsetSize(metadata);
+ int numKeys = getMetadataNumKeys(metadata, offsetSize);
+ int stringStart = getMetadataStringStart(metadata, numKeys, offsetSize);
+ int lastOffset = readUnsigned(metadata, stringStart - offsetSize,
offsetSize);
+ if (lastOffset > metadata.length - stringStart) {
Review Comment:
Is it allowed for `metadata.length > stringStart + lastOffset`? Probably not
right, because you presumably want to compare raw binaries?
##########
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;
Review Comment:
If a new type is introduced into the spec which requires additional
canonicalization, this function will break. It is hard to keep track of every
single line of code to update when you introduce something new.
Can we handle every type explicitly and have the default be false?
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -282,23 +358,25 @@ public void appendBoolean(boolean b) {
// range of the long value.
public void appendLong(long l) {
checkCapacity(1 + 8);
- if (l == (byte) l) {
- writeBuffer[writePos++] = primitiveHeader(INT1);
- writeLong(writeBuffer, writePos, l, 1);
- writePos += 1;
- } else if (l == (short) l) {
- writeBuffer[writePos++] = primitiveHeader(INT2);
- writeLong(writeBuffer, writePos, l, 2);
- writePos += 2;
- } else if (l == (int) l) {
- writeBuffer[writePos++] = primitiveHeader(INT4);
- writeLong(writeBuffer, writePos, l, 4);
- writePos += 4;
- } else {
- writeBuffer[writePos++] = primitiveHeader(INT8);
- writeLong(writeBuffer, writePos, l, 8);
- writePos += 8;
+ int typeInfo = minimalLongTypeInfo(l);
Review Comment:
We are refactoring a lot of code that is not related to the canonicalization
effort. This makes sense given that we now have new APIs but would it be better
to have that be handled as a follow up? My main concern is that
canonicalization is a major feature for whom it would be nice to have a single
commit show all the necessary code.
It is hard for me to stamp this PR for the canonicalization changes as well
as meticulously making sure that no existing code is broken or regresses in
performance.
##########
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:
The private method is only used once. Can we inline it? Doesn't seem much
more complicated than the decimal logic
##########
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) {
+ int encodedSize = valueSize(value, pos);
+ if (encodedSize <= 0 || (long) pos + encodedSize > value.length) {
Review Comment:
`isValueCanonical` already performs this check. You'd be able to shave this
code of if you inline 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]
