[SPARK-48175][SQL][PYTHON] Store collation information in metadata and not in type for SER/DE

common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java

@@ -36,11 +36,42 @@
  * Provides functionality to the UTF8String object which respects defined collation settings.
  */
 public final class CollationFactory {
+
+ public static class CollationIdentifier {
+ public final String provider;
+ public final String name;
+ public final String version;
+
+ public CollationIdentifier(String provider, String collationName, String version) {
+ this.provider = provider;
+ this.name = collationName;
+ this.version = version;
+ }
+
+ public static CollationIdentifier fromString(String identifier) {
+ String[] parts = identifier.split("\\.", 3);
+ return new CollationIdentifier(parts[0], parts[1], parts[2]);
+ }
+
+ @Override
+ public String toString() {
+ return String.format("%s.%s.%s", provider, name, version);
+ }
+
+ /**
+ * Returns
+ */
+ public String versionLess() {
+ return String.format("%s.%s", provider, name);
+ }
+ }
+
  /**
  * Entry encapsulating all information about a collation.
  */
  public static class Collation {
  public final String collationName;
+ public final String provider;
  public final Collator collator;
  public final Comparator<UTF8String> comparator;
 
@@ -97,6 +128,7 @@ public Collation(
  boolean supportsBinaryOrdering,
  boolean supportsLowercaseEquality) {
  this.collationName = collationName;
+ this.provider = collator == null ? "spark" : "icu";
  this.collator = collator;
  this.comparator = comparator;
  this.version = version;
@@ -137,6 +169,11 @@ public Collation(
  supportsBinaryOrdering,
  supportsLowercaseEquality);
  }
+
+ /** Returns the collation identifier*/
+ public CollationIdentifier identifier() {
+ return new CollationIdentifier(provider, collationName, version);
+ }
  }
 
  private static final Collation[] collationTable = new Collation[4];

common/utils/src/main/resources/error/error-conditions.json

@@ -2322,6 +2322,12 @@
  ],
  "sqlState" : "2203G"
  },
+ "INVALID_JSON_DATA_TYPE_FOR_COLLATIONS" : {
+ "message" : [
+ "Collations can only be applied to string types, but the JSON data type is <jsonType>."
+ ],
+ "sqlState" : "2203G"
+ },
  "INVALID_JSON_ROOT_FIELD" : {
  "message" : [
  "Cannot convert JSON root field to target Spark type."

python/pyspark/errors/error-conditions.json

@@ -357,6 +357,11 @@
  "All items in `<arg_name>` should be in <allowed_types>, got <item_type>."
  ]
  },
+ "INVALID_JSON_DATA_TYPE_FOR_COLLATIONS" : {
+ "message" : [
+ "Collations can only be applied to string types, but the JSON data type is <jsonType>."
+ ]
+ },
  "INVALID_MULTIPLE_ARGUMENT_CONDITIONS": {
  "message": [
  "[{arg_names}] cannot be <condition>."

python/pyspark/sql/tests/connect/test_parity_types.py

@@ -86,6 +86,10 @@ def test_rdd_with_udt(self):
  def test_udt(self):
  super().test_udt()
 
+ @unittest.skip("Requires JVM access.")
+ def test_schema_with_collations_json_ser_de(self):
+ super().test_schema_with_collations_json_ser_de()
+
  @unittest.skip("Does not test anything related to Spark Connect")
  def test_parse_datatype_string(self):
  super().test_parse_datatype_string()

python/pyspark/sql/tests/test_types.py

@@ -549,6 +549,129 @@ def test_convert_list_to_str(self):
  self.assertEqual(df.count(), 1)
  self.assertEqual(df.head(), Row(name="[123]", income=120))
 
+ def test_schema_with_collations_json_ser_de(self):
+ from pyspark.sql.types import _parse_datatype_json_string
+
+ unicode_collation = "UNICODE"
+
+ simple_struct = StructType([StructField("c1", StringType(unicode_collation))])
+
+ nested_struct = StructType([StructField("nested", simple_struct)])
+
+ array_in_schema = StructType(
+ [StructField("array", ArrayType(StringType(unicode_collation)))]
+ )
+
+ map_in_schema = StructType(
+ [
+ StructField(
+ "map", MapType(StringType(unicode_collation), StringType(unicode_collation))
+ )
+ ]
+ )
+
+ array_in_map = StructType(
+ [
+ StructField(
+ "arrInMap",
+ MapType(
+ StringType(unicode_collation), ArrayType(StringType(unicode_collation))
+ ),
+ )
+ ]
+ )
+
+ nested_array_in_map = StructType(
+ [
+ StructField(
+ "nestedArrayInMap",
+ ArrayType(
+ MapType(
+ StringType(unicode_collation),
+ ArrayType(ArrayType(StringType(unicode_collation))),
+ )
+ ),
+ )
+ ]
+ )
+
+ schema_with_multiple_fields = StructType(
+ simple_struct.fields
+ + nested_struct.fields
+ + array_in_schema.fields
+ + map_in_schema.fields
+ + array_in_map.fields
+ + nested_array_in_map.fields
+ )
+
+ schemas = [
+ simple_struct,
+ nested_struct,
+ array_in_schema,
+ map_in_schema,
+ nested_array_in_map,
+ array_in_map,
+ schema_with_multiple_fields,
+ ]
+
+ for schema in schemas:
+ scala_datatype = self.spark._jsparkSession.parseDataType(schema.json())
+ python_datatype = _parse_datatype_json_string(scala_datatype.json())
+ assert schema == python_datatype
+ assert schema == _parse_datatype_json_string(schema.json())
+
+ def test_schema_with_collations_on_non_string_types(self):
+ from pyspark.sql.types import _parse_datatype_json_string, _COLLATIONS_METADATA_KEY
+
+ collations_on_int_col_json = f"""
+ {{
+ "type": "struct",
+ "fields": [
+ {{
+ "name": "c1",
+ "type": "integer",
+ "nullable": true,
+ "metadata": {{
+ "{_COLLATIONS_METADATA_KEY}": {{
+ "c1": "icu.UNICODE"
+ }}
+ }}
+ }}
+ ]
+ }}
+ """
+
+ collations_in_map_json = f"""
+ {{
+ "type": "struct",
+ "fields": [
+ {{
+ "name": "mapField",
+ "type": {{
+ "type": "map",
+ "keyType": "string",
+ "valueType": "integer",
+ "valueContainsNull": true
+ }},
+ "nullable": true,
+ "metadata": {{
+ "{_COLLATIONS_METADATA_KEY}": {{
+ "mapField.value": "icu.UNICODE"
+ }}
+ }}
+ }}
+ ]
+ }}
+ """
+
+ self.assertRaises(
+ PySparkTypeError, lambda: _parse_datatype_json_string(collations_on_int_col_json)
+ )
+
+ self.assertRaises(
+ PySparkTypeError, lambda: _parse_datatype_json_string(collations_in_map_json)
+ )
+
  def test_udt(self):
  from pyspark.sql.types import _parse_datatype_json_string, _infer_type, _make_type_verifier
 
@@ -864,27 +987,6 @@ def test_parse_datatype_string(self):
  self.assertEqual(t(), _parse_datatype_string(k))
  self.assertEqual(IntegerType(), _parse_datatype_string("int"))
  self.assertEqual(StringType(), _parse_datatype_string("string"))
- self.assertEqual(StringType(), _parse_datatype_string("string collate UTF8_BINARY"))
- self.assertEqual(StringType(), _parse_datatype_string("string COLLATE UTF8_BINARY"))
- self.assertEqual(
- StringType.fromCollationId(0), _parse_datatype_string("string COLLATE UTF8_BINARY")
- )
- self.assertEqual(
- StringType.fromCollationId(1),
- _parse_datatype_string("string COLLATE UTF8_BINARY_LCASE"),
- )
- self.assertEqual(
- StringType.fromCollationId(2), _parse_datatype_string("string COLLATE UNICODE")
- )
- self.assertEqual(
- StringType.fromCollationId(2), _parse_datatype_string("string COLLATE `UNICODE`")
- )
- self.assertEqual(
- StringType.fromCollationId(3), _parse_datatype_string("string COLLATE UNICODE_CI")
- )
- self.assertEqual(
- StringType.fromCollationId(3), _parse_datatype_string("string COLLATE `UNICODE_CI`")
- )
  self.assertEqual(CharType(1), _parse_datatype_string("char(1)"))
  self.assertEqual(CharType(10), _parse_datatype_string("char( 10 )"))
  self.assertEqual(CharType(11), _parse_datatype_string("char( 11)"))