[Kernel] Extended schema JSON serde to support collations

kernel/kernel-api/src/main/java/io/delta/kernel/internal/types/DataTypeJsonSerDe.java

@@ -89,7 +89,11 @@ public static String serializeDataType(DataType dataType) {
    */
   public static StructType deserializeStructType(String structTypeJson) {
     try {
-      DataType parsedType = parseDataType(OBJECT_MAPPER.reader().readTree(structTypeJson));
+      DataType parsedType =
+          parseDataType(
+              OBJECT_MAPPER.reader().readTree(structTypeJson),
+              "" /* fieldPath */,
+              new FieldMetadata.Builder().build() /* collationsMetadata */);
       if (parsedType instanceof StructType) {
         return (StructType) parsedType;
       } else {
@@ -130,23 +134,61 @@ public static StructType deserializeStructType(String structTypeJson) {
    *   "nullable" : false,
    *   "metadata" : { }
    * }
+   *
+   * // Collated string type field serialized as:
+   * {
+   *   "name" : "s",
+   *   "type" : "string",
+   *   "nullable", false,
+   *   "metadata" : {
+   *     "__COLLATIONS": { "s": "ICU.de_DE" }
+   *   }
+   * }
+   *
+   * // Array with collated strings field serialized as:
+   * {
+   *   "name" : "arr",
+   *   "type" : {
+   *     "type" : "array",
+   *     "elementType" : "string",
+   *     "containsNull" : false
+   *   }
+   *   "nullable" : false,
+   *   "metadata" : {
+   *     "__COLLATIONS": { "arr.element": "ICU.de_DE" }
+   *   }
+   * }
    * </pre>
+   *
+   * @param fieldPath Path from the nearest ancestor that is of the {@link StructField} type. For
+   *     example, "c1.key.element" represents a path starting from the {@link StructField} named
+   *     "c1." The next element, "key," indicates that "c1" stores a {@link MapType} type. The final
+   *     element, "element", shows that the key of the map is an {@link ArrayType} type.
+   * @param collationsMetadata Metadata that maps the path of a {@link StringType} to its collation.
+   *     Only maps non-UTF8_BINARY collated {@link StringType}. Collation metadata is stored in the
+   *     nearest ancestor, which is the StructField. This is because StructField includes a metadata
+   *     field, whereas Map and Array do not, making them unable to store this information. Paths
+   *     are in same form as `fieldPath`. <a
+   *     href="https://github.com/delta-io/delta/blob/master/protocol_rfcs/collated-string-type.md#collation-identifiers">Docs</a>
    */
-  static DataType parseDataType(JsonNode json) {
+  static DataType parseDataType(JsonNode json, String fieldPath, FieldMetadata collationsMetadata) {
     switch (json.getNodeType()) {
       case STRING:
         // simple types are stored as just a string
-        return nameToType(json.textValue());
+        return nameToType(json.textValue(), fieldPath, collationsMetadata);
       case OBJECT:
         // complex types (array, map, or struct are stored as JSON objects)
         String type = getStringField(json, "type");
         switch (type) {
           case "struct":
+            assertValidTypeForCollations(fieldPath, "struct", collationsMetadata);
             return parseStructType(json);
           case "array":
-            return parseArrayType(json);
+            assertValidTypeForCollations(fieldPath, "array", collationsMetadata);
+            return parseArrayType(json, fieldPath, collationsMetadata);
           case "map":
-            return parseMapType(json);
+            assertValidTypeForCollations(fieldPath, "map", collationsMetadata);
+            return parseMapType(json, fieldPath, collationsMetadata);
             // No default case here; fall through to the following error when no match
         }
       default:
@@ -160,26 +202,32 @@ static DataType parseDataType(JsonNode json) {
    * Parses an <a href="https://github.com/delta-io/delta/blob/master/PROTOCOL.md#array-type">array
    * type </a>
    */
-  private static ArrayType parseArrayType(JsonNode json) {
+  private static ArrayType parseArrayType(
+      JsonNode json, String fieldPath, FieldMetadata collationsMetadata) {
     checkArgument(
         json.isObject() && json.size() == 3,
         String.format("Expected JSON object with 3 fields for array data type but got:\n%s", json));
     boolean containsNull = getBooleanField(json, "containsNull");
-    DataType dataType = parseDataType(getNonNullField(json, "elementType"));
+    DataType dataType =
+        parseDataType(
+            getNonNullField(json, "elementType"), fieldPath + ".element", collationsMetadata);
     return new ArrayType(dataType, containsNull);
   }
 
   /**
    * Parses an <a href="https://github.com/delta-io/delta/blob/master/PROTOCOL.md#map-type">map type
    * </a>
    */
-  private static MapType parseMapType(JsonNode json) {
+  private static MapType parseMapType(
+      JsonNode json, String fieldPath, FieldMetadata collationsMetadata) {
     checkArgument(
         json.isObject() && json.size() == 4,
         String.format("Expected JSON object with 4 fields for map data type but got:\n%s", json));
     boolean valueContainsNull = getBooleanField(json, "valueContainsNull");
-    DataType keyType = parseDataType(getNonNullField(json, "keyType"));
-    DataType valueType = parseDataType(getNonNullField(json, "valueType"));
+    DataType keyType =
+        parseDataType(getNonNullField(json, "keyType"), fieldPath + ".key", collationsMetadata);
+    DataType valueType =
+        parseDataType(getNonNullField(json, "valueType"), fieldPath + ".value", collationsMetadata);
     return new MapType(keyType, valueType, valueContainsNull);
   }
 
@@ -211,14 +259,25 @@ private static StructType parseStructType(JsonNode json) {
   private static StructField parseStructField(JsonNode json) {
     Preconditions.checkArgument(json.isObject(), "Expected JSON object for struct field");
     String name = getStringField(json, "name");
-    DataType type = parseDataType(getNonNullField(json, "type"));
+    FieldMetadata metadata = parseFieldMetadata(json.get("metadata"), false);
+    DataType type =
+        parseDataType(
+            getNonNullField(json, "type"), name, getCollationsMetadata(json.get("metadata")));
     boolean nullable = getBooleanField(json, "nullable");
-    FieldMetadata metadata = parseFieldMetadata(json.get("metadata"));
     return new StructField(name, type, nullable, metadata);
   }
 
   /** Parses an {@link FieldMetadata}. */
   private static FieldMetadata parseFieldMetadata(JsonNode json) {
+    return parseFieldMetadata(json, true);
+  }
+
+  /**
+   * Parses a {@link FieldMetadata}, optionally including collation metadata, depending on
+   * `includecollationsMetadata`.
+   */
+  private static FieldMetadata parseFieldMetadata(
+      JsonNode json, boolean includecollationsMetadata) {
     if (json == null || json.isNull()) {
       return FieldMetadata.empty();
     }
@@ -231,6 +290,10 @@ private static FieldMetadata parseFieldMetadata(JsonNode json) {
       JsonNode value = entry.getValue();
       String key = entry.getKey();
 
+      if (!includecollationsMetadata && key.equals(DataType.COLLATIONS_METADATA_KEY)) {
+        continue;
+      }
+
       if (value.isNull()) {
         builder.putNull(key);
       } else if (value.isIntegralNumber()) { // covers both int and long
@@ -298,8 +361,13 @@ private static <T> List<T> buildList(JsonNode json, Function<JsonNode, T> access
   private static Pattern FIXED_DECIMAL_PATTERN = Pattern.compile(FIXED_DECIMAL_REGEX);
 
   /** Parses primitive string type names to a {@link DataType} */
-  private static DataType nameToType(String name) {
+  private static DataType nameToType(
+      String name, String fieldPath, FieldMetadata collationsMetadata) {
     if (BasePrimitiveType.isPrimitiveType(name)) {
+      if (collationsMetadata.contains(fieldPath)) {
+        assertValidTypeForCollations(fieldPath, name, collationsMetadata);
+        return new StringType(collationsMetadata.getString(fieldPath));
+      }
       return BasePrimitiveType.createPrimitive(name);
     } else if (name.equals("decimal")) {
       return DecimalType.USER_DEFAULT;
@@ -341,6 +409,22 @@ private static String getStringField(JsonNode rootNode, String fieldName) {
     return node.textValue(); // double check this only works for string values! and isTextual()!
   }
 
+  private static void assertValidTypeForCollations(
+      String fieldPath, String fieldType, FieldMetadata collationsMetadata) {
+    if (collationsMetadata.contains(fieldPath) && !fieldType.equals("string")) {
+      throw new IllegalArgumentException(
+          String.format("Invalid data type for collations: \"%s\"", fieldType));
+    }
+  }
+
+  /** Returns a metadata with a map of field path to collation name. */
+  private static FieldMetadata getCollationsMetadata(JsonNode fieldMetadata) {
+    if (fieldMetadata == null || !fieldMetadata.has(DataType.COLLATIONS_METADATA_KEY)) {
+      return new FieldMetadata.Builder().build();
+    }
+    return parseFieldMetadata(fieldMetadata.get(DataType.COLLATIONS_METADATA_KEY));
+  }
+
   private static boolean getBooleanField(JsonNode rootNode, String fieldName) {
     JsonNode node = getNonNullField(rootNode, fieldName);
     Preconditions.checkArgument(
@@ -414,7 +498,7 @@ private static void writeStructField(JsonGenerator gen, StructField field) throw
     writeDataType(gen, field.getDataType());
     gen.writeBooleanField("nullable", field.isNullable());
     gen.writeFieldName("metadata");
-    writeFieldMetadata(gen, field.getMetadata());
+    writeFieldMetadata(gen, field.getSerializationMetadata());
     gen.writeEndObject();
   }
 

kernel/kernel-api/src/main/java/io/delta/kernel/types/DataType.java

@@ -25,6 +25,7 @@
  */
 @Evolving
 public abstract class DataType {
+  public static final String COLLATIONS_METADATA_KEY = "__COLLATIONS";
 
   /**
    * Are the data types same? The metadata or column names could be different.

kernel/kernel-api/src/main/java/io/delta/kernel/types/StructField.java

@@ -17,6 +17,9 @@
 package io.delta.kernel.types;
 
 import io.delta.kernel.annotation.Evolving;
+import io.delta.kernel.internal.util.Tuple2;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 
 /**
@@ -102,6 +105,47 @@ public String toString() {
         "StructField(name=%s,type=%s,nullable=%s,metadata=%s)", name, dataType, nullable, metadata);
   }
 
+  public FieldMetadata getSerializationMetadata() {
+    List<Tuple2<String, String>> nestedCollatedFields = getNestedCollatedFields(dataType, name);
+    if (nestedCollatedFields.isEmpty()) {
+      return metadata;
+    }
+
+    FieldMetadata.Builder metadataBuilder = new FieldMetadata.Builder();
+    for (Tuple2<String, String> nestedField : nestedCollatedFields) {
+      metadataBuilder.putString(nestedField._1, nestedField._2);
+    }
+    return new FieldMetadata.Builder()
+        .fromMetadata(metadata)
+        .putFieldMetadata(DataType.COLLATIONS_METADATA_KEY, metadataBuilder.build())
+        .build();
+  }
+
+  private List<Tuple2<String, String>> getNestedCollatedFields(DataType parent, String path) {
+    List<Tuple2<String, String>> nestedCollatedFields = new ArrayList<>();
+    if (parent instanceof StringType) {
+      StringType stringType = (StringType) parent;
+      if (!stringType
+          .getCollationIdentifier()
+          .equals(CollationIdentifier.fromString("SPARK.UTF8_BINARY"))) {
+        nestedCollatedFields.add(
+            new Tuple2<>(
+                path, ((StringType) parent).getCollationIdentifier().toStringWithoutVersion()));
+      }
+    } else if (parent instanceof MapType) {
+      nestedCollatedFields.addAll(
+          getNestedCollatedFields(((MapType) parent).getKeyType(), path + ".key"));
+      nestedCollatedFields.addAll(
+          getNestedCollatedFields(((MapType) parent).getValueType(), path + ".value"));
+    } else if (parent instanceof ArrayType) {
+      nestedCollatedFields.addAll(
+          getNestedCollatedFields(((ArrayType) parent).getElementType(), path + ".element"));
+    }
+    // We didn't check for StructType because we store the StringType's
+    // collation information in the nearest ancestor StructField's metadata when serializing.
+    return nestedCollatedFields;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {