SavicStefan commented on code in PR #59669:
URL: https://github.com/apache/spark/pull/59669#discussion_r4386841891


##########
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:
   So if we choose to preserve the bytes without sanitizing invalid UTF-8, then 
plain GROUP BY would not throw for invalid UTF-8, if we want it to fail there 
we need to validate it during canonicalization and throw there.
   
   I am leaning more towards preserving the original bytes, because sanitizing 
can turn different invalid byte sequences into the same replacement character 
and group them together
   
   wdyt?



-- 
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]

Reply via email to