dbatomic commented on code in PR #46180:
URL: https://github.com/apache/spark/pull/46180#discussion_r1576145404
##########
common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java:
##########
@@ -117,76 +118,490 @@ public Collation(
}
/**
- * Constructor with comparators that are inherited from the given collator.
+ * collation id (32-bit integer) layout:
+ * bit 31: 0 = predefined collation, 1 = user-defined collation
+ * bit 30: 0 = utf8-binary, 1 = ICU
+ * bit 29: 0 = case-sensitive, 1 = case-insensitive
+ * bit 28: 0 = accent-sensitive, 1 = accent-insensitive
+ * bit 27-26: 00 = unspecified, 01 = punctuation-sensitive, 10 =
punctuation-insensitive
+ * bit 25-24: 00 = unspecified, 01 = first-lower, 10 = first-upper
+ * bit 23-22: 00 = unspecified, 01 = to-lower, 10 = to-upper
+ * bit 21-20: 00 = unspecified, 01 = trim-left, 10 = trim-right, 11 =
trim-both
+ * bit 19-18: zeroes, reserved for version
+ * bit 17-16: zeroes
+ * bit 15-0: locale id for ICU collations / zeroes for utf8-binary
*/
- public Collation(
- String collationName,
- Collator collator,
- String version,
- boolean supportsBinaryEquality,
- boolean supportsBinaryOrdering,
- boolean supportsLowercaseEquality) {
- this(
- collationName,
- collator,
- (s1, s2) -> collator.compare(s1.toString(), s2.toString()),
- version,
- s -> (long)collator.getCollationKey(s.toString()).hashCode(),
- supportsBinaryEquality,
- supportsBinaryOrdering,
- supportsLowercaseEquality);
- }
- }
+ private static class CollationSpec {
+ private enum ImplementationProvider {
+ UTF8_BINARY, ICU
+ }
+
+ private enum CaseSensitivity {
+ CS, CI
+ }
+
+ private enum AccentSensitivity {
+ AS, AI
+ }
+
+ private enum PunctuationSensitivity {
+ UNSPECIFIED, PS, PI
+ }
- private static final Collation[] collationTable = new Collation[4];
- private static final HashMap<String, Integer> collationNameToIdMap = new
HashMap<>();
-
- public static final int UTF8_BINARY_COLLATION_ID = 0;
- public static final int UTF8_BINARY_LCASE_COLLATION_ID = 1;
-
- static {
- // Binary comparison. This is the default collation.
- // No custom comparators will be used for this collation.
- // Instead, we rely on byte for byte comparison.
- collationTable[0] = new Collation(
- "UTF8_BINARY",
- null,
- UTF8String::binaryCompare,
- "1.0",
- s -> (long)s.hashCode(),
- true,
- true,
- false);
-
- // Case-insensitive UTF8 binary collation.
- // TODO: Do in place comparisons instead of creating new strings.
- collationTable[1] = new Collation(
- "UTF8_BINARY_LCASE",
- null,
- UTF8String::compareLowerCase,
- "1.0",
- (s) -> (long)s.toLowerCase().hashCode(),
- false,
- false,
- true);
-
- // UNICODE case sensitive comparison (ROOT locale, in ICU).
- collationTable[2] = new Collation(
- "UNICODE", Collator.getInstance(ULocale.ROOT), "153.120.0.0", true,
false, false);
- collationTable[2].collator.setStrength(Collator.TERTIARY);
- collationTable[2].collator.freeze();
-
- // UNICODE case-insensitive comparison (ROOT locale, in ICU + Secondary
strength).
- collationTable[3] = new Collation(
- "UNICODE_CI", Collator.getInstance(ULocale.ROOT), "153.120.0.0", false,
false, false);
- collationTable[3].collator.setStrength(Collator.SECONDARY);
- collationTable[3].collator.freeze();
-
- for (int i = 0; i < collationTable.length; i++) {
- collationNameToIdMap.put(collationTable[i].collationName, i);
+ private enum FirstLetterPreference {
+ UNSPECIFIED, FU, FL
+ }
+
+ private enum CaseConversion {
+ UNSPECIFIED, LCASE, UCASE
+ }
+
+ private enum SpaceTrimming {
+ UNSPECIFIED, LTRIM, RTRIM, TRIM
+ }
+
+ private static final int implementationProviderOffset = 30;
+ private static final int implementationProviderLen = 1;
+ private static final int caseSensitivityOffset = 29;
+ private static final int caseSensitivityLen = 1;
+ private static final int accentSensitivityOffset = 28;
+ private static final int accentSensitivityLen = 1;
+ private static final int punctuationSensitivityOffset = 26;
+ private static final int punctuationSensitivityLen = 2;
+ private static final int firstLetterPreferenceOffset = 24;
+ private static final int firstLetterPreferenceLen = 2;
+ private static final int caseConversionOffset = 22;
+ private static final int caseConversionLen = 2;
+ private static final int spaceTrimmingOffset = 20;
+ private static final int spaceTrimmingLen = 2;
+ private static final int localeOffset = 0;
+ private static final int localeLen = 16;
+
+ private static final String[] ICULocaleNames;
+ private static final Map<String, ULocale> ICULocaleMap = new HashMap<>();
+ private static final Map<String, String> ICULocaleMapUppercase = new
HashMap<>();
+ private static final Map<String, Integer> ICULocaleToId = new
HashMap<>();
+
+ static {
+ ICULocaleMap.put("UNICODE", ULocale.ROOT);
+ ULocale[] locales = Collator.getAvailableULocales();
+ for (ULocale locale : locales) {
+ if (locale.getVariant().isEmpty()) {
+ String language = locale.getLanguage();
+ assert (!language.isEmpty());
+ StringBuilder builder = new StringBuilder(language);
+ String script = locale.getScript();
+ if (!script.isEmpty()) {
+ builder.append('_');
+ builder.append(script);
+ }
+ String country = locale.getISO3Country();
+ if (!country.isEmpty()) {
+ builder.append('_');
+ builder.append(country);
+ }
+ String localeName = builder.toString();
+ assert (!ICULocaleMap.containsKey(localeName));
+ ICULocaleMap.put(localeName, locale);
+ }
+ }
+ for (String localeName : ICULocaleMap.keySet()) {
+ String localeUppercase = localeName.toUpperCase();
+ assert (!ICULocaleMapUppercase.containsKey(localeUppercase));
+ ICULocaleMapUppercase.put(localeUppercase, localeName);
+ }
+ ICULocaleNames = ICULocaleMap.keySet().toArray(new String[0]);
+ Arrays.sort(ICULocaleNames);
+ assert (ICULocaleNames.length <= (1 << 16));
+ for (int i = 0; i < ICULocaleNames.length; i++) {
+ ICULocaleToId.put(ICULocaleNames[i], i);
+ }
+ }
+
+ public static final int UTF8_BINARY_COLLATION_ID =
+ new CollationSpec(
+ ImplementationProvider.UTF8_BINARY,
+ null,
+ CaseSensitivity.CS,
+ AccentSensitivity.AS,
+ PunctuationSensitivity.UNSPECIFIED,
+ FirstLetterPreference.UNSPECIFIED,
+ CaseConversion.UNSPECIFIED,
+ SpaceTrimming.UNSPECIFIED
+ ).getCollationId();
+
+ public static final int UTF8_BINARY_LCASE_COLLATION_ID =
+ new CollationSpec(
+ ImplementationProvider.UTF8_BINARY,
+ null,
+ CaseSensitivity.CS,
+ AccentSensitivity.AS,
+ PunctuationSensitivity.UNSPECIFIED,
+ FirstLetterPreference.UNSPECIFIED,
+ CaseConversion.LCASE,
+ SpaceTrimming.UNSPECIFIED
+ ).getCollationId();
+
+ public static final int UNICODE_COLLATION_ID =
+ new CollationSpec(
+ ImplementationProvider.ICU,
+ "UNICODE",
+ CaseSensitivity.CS,
+ AccentSensitivity.AS,
+ PunctuationSensitivity.UNSPECIFIED,
+ FirstLetterPreference.UNSPECIFIED,
+ CaseConversion.UNSPECIFIED,
+ SpaceTrimming.UNSPECIFIED
+ ).getCollationId();
+
+ public static final int UNICODE_CI_COLLATION_ID =
+ new CollationSpec(
+ ImplementationProvider.ICU,
+ "UNICODE",
+ CaseSensitivity.CI,
+ AccentSensitivity.AS,
+ PunctuationSensitivity.UNSPECIFIED,
+ FirstLetterPreference.UNSPECIFIED,
+ CaseConversion.UNSPECIFIED,
+ SpaceTrimming.UNSPECIFIED
+ ).getCollationId();
+
+ private final ImplementationProvider implementationProvider;
+ private final CaseSensitivity caseSensitivity;
+ private final AccentSensitivity accentSensitivity;
+ private final PunctuationSensitivity punctuationSensitivity;
+ private final FirstLetterPreference firstLetterPreference;
+ private final CaseConversion caseConversion;
+ private final SpaceTrimming spaceTrimming;
+ private final String locale;
+ private final int collationId;
+
+ private CollationSpec(
+ ImplementationProvider implementationProvider,
+ String locale,
+ CaseSensitivity caseSensitivity,
+ AccentSensitivity accentSensitivity,
+ PunctuationSensitivity punctuationSensitivity,
+ FirstLetterPreference firstLetterPreference,
+ CaseConversion caseConversion,
+ SpaceTrimming spaceTrimming) {
+ this.implementationProvider = implementationProvider;
+ this.locale = locale;
+ this.caseSensitivity = caseSensitivity;
+ this.accentSensitivity = accentSensitivity;
+ this.punctuationSensitivity = punctuationSensitivity;
+ this.firstLetterPreference = firstLetterPreference;
+ this.caseConversion = caseConversion;
+ this.spaceTrimming = spaceTrimming;
+ this.collationId = getCollationId();
+ }
+
+ public int getCollationId() {
+ int collationId = 0;
+ collationId |= implementationProvider.ordinal() <<
implementationProviderOffset;
+ collationId |= caseSensitivity.ordinal() << caseSensitivityOffset;
+ collationId |= accentSensitivity.ordinal() << accentSensitivityOffset;
+ collationId |= punctuationSensitivity.ordinal() <<
punctuationSensitivityOffset;
+ collationId |= firstLetterPreference.ordinal() <<
firstLetterPreferenceOffset;
+ collationId |= caseConversion.ordinal() << caseConversionOffset;
+ collationId |= spaceTrimming.ordinal() << spaceTrimmingOffset;
+ if (implementationProvider == ImplementationProvider.ICU) {
+ collationId |= ICULocaleToId.get(locale);
+ }
+ return collationId;
+ }
+
+ public static CollationSpec fromCollationId(int collationId) {
+ ImplementationProvider implementationProvider =
ImplementationProvider.values()[
+ (collationId >> implementationProviderOffset) & ((1 <<
implementationProviderLen) - 1)];
+ CaseSensitivity caseSensitivity = CaseSensitivity.values()[
+ (collationId >> caseSensitivityOffset) & ((1 << caseSensitivityLen)
- 1)];
+ AccentSensitivity accentSensitivity = AccentSensitivity.values()[
+ (collationId >> accentSensitivityOffset) & ((1 <<
accentSensitivityLen) - 1)];
+ PunctuationSensitivity punctuationSensitivity =
PunctuationSensitivity.values()[
+ (collationId >> punctuationSensitivityOffset) & ((1 <<
punctuationSensitivityLen) - 1)];
+ FirstLetterPreference firstLetterPreference =
FirstLetterPreference.values()[
+ (collationId >> firstLetterPreferenceOffset) & ((1 <<
firstLetterPreferenceLen) - 1)];
+ CaseConversion caseConversion = CaseConversion.values()[
+ (collationId >> caseConversionOffset) & ((1 << caseConversionLen) -
1)];
+ SpaceTrimming spaceTrimming = SpaceTrimming.values()[
+ (collationId >> spaceTrimmingOffset) & ((1 << spaceTrimmingLen) -
1)];
+ String locale;
+ if (implementationProvider == ImplementationProvider.UTF8_BINARY) {
+ locale = "UTF8_BINARY";
+ } else {
+ locale = ICULocaleNames[(collationId >> localeOffset) & ((1 <<
localeLen) - 1)];
+ }
+ return new CollationSpec(
+ implementationProvider,
+ locale,
+ caseSensitivity,
+ accentSensitivity,
+ punctuationSensitivity,
+ firstLetterPreference,
+ caseConversion,
+ spaceTrimming
+ );
+ }
+
+ public static int collationNameToId(String originalCollationName) throws
SparkException {
+ String collationName = originalCollationName.toUpperCase();
+ try {
+ if (collationName.startsWith("UTF8_BINARY")) {
+ return collationUTF8BinaryNameToId(collationName);
+ } else {
+ return collationICUNameToId(collationName);
+ }
+ } catch (SparkException e) {
+ throw new SparkException(
+ "COLLATION_INVALID_NAME",
+ SparkException.constructMessageParams(Map.of("collationName",
originalCollationName)),
+ e);
+ }
+ }
+
+ private static int collationUTF8BinaryNameToId(String collationName)
throws SparkException {
+ int collationId = 0;
+ collationId |= ImplementationProvider.UTF8_BINARY.ordinal() <<
implementationProviderOffset;
+ collationId |=
parseSpecifiers(collationName.substring("UTF8_BINARY".length()));
+ return collationId;
+ }
+
+ private static int collationICUNameToId(String collationName) throws
SparkException {
+ int lastPos = -1;
+ for (int i = 1; i <= collationName.length(); i++) {
+ String localeName = collationName.substring(0, i);
+ if (ICULocaleMapUppercase.containsKey(localeName)) {
+ lastPos = i;
+ }
+ }
+ if (lastPos == -1) {
+ throw new SparkException("Invalid locale in collation name value " +
collationName);
+ } else {
+ int collationId = 0;
+ collationId |= ImplementationProvider.ICU.ordinal() <<
implementationProviderOffset;
+ collationId |= parseSpecifiers(collationName.substring(lastPos));
+ String normalizedLocaleName = ICULocaleMapUppercase.get(
+ collationName.substring(0, lastPos));
+ collationId |= ICULocaleToId.get(normalizedLocaleName) <<
localeOffset;
+ return collationId;
+ }
+ }
+
+ private static int parseSpecifiers(String specString) throws
SparkException {
+ int specifiers = 0;
+ String[] parts = specString.split("_");
+ for (String part : parts) {
+ if (!part.isEmpty()) {
+ if (part.equals("UNSPECIFIED")) {
+ throw new SparkException("UNSPECIFIED collation specifier
reserved for internal use");
+ } else if (Arrays.stream(CaseSensitivity.values()).anyMatch(
+ (s) -> s.toString().equals(part))) {
+ specifiers |=
+ CaseSensitivity.valueOf(part).ordinal() <<
caseSensitivityOffset;
+ } else if (Arrays.stream(AccentSensitivity.values()).anyMatch(
+ (s) -> part.equals(s.toString()))) {
+ specifiers |=
+ AccentSensitivity.valueOf(part).ordinal() <<
accentSensitivityOffset;
+ } else if (Arrays.stream(PunctuationSensitivity.values()).anyMatch(
+ (s) -> part.equals(s.toString()))) {
+ specifiers |=
+ PunctuationSensitivity.valueOf(part).ordinal() <<
punctuationSensitivityOffset;
+ } else if (Arrays.stream(FirstLetterPreference.values()).anyMatch(
+ (s) -> part.equals(s.toString()))) {
+ specifiers |=
+ FirstLetterPreference.valueOf(part).ordinal() <<
firstLetterPreferenceOffset;
+ } else if (Arrays.stream(CaseConversion.values()).anyMatch(
+ (s) -> part.equals(s.toString()))) {
+ specifiers |=
+ CaseConversion.valueOf(part).ordinal() << caseConversionOffset;
+ } else if (Arrays.stream(SpaceTrimming.values()).anyMatch(
+ (s) -> part.equals(s.toString()))) {
+ specifiers |=
+ SpaceTrimming.valueOf(part).ordinal() << spaceTrimmingOffset;
+ } else {
+ throw new SparkException("Invalid collation specifier value " +
part);
+ }
+ }
+ }
+ return specifiers;
+ }
+
+ public Collation buildCollation() {
+ if (implementationProvider == ImplementationProvider.UTF8_BINARY) {
+ return buildUTF8BinaryCollation();
+ } else {
+ return buildICUCollation();
+ }
+ }
+
+ public Collation buildUTF8BinaryCollation() {
+ Comparator<UTF8String> comparator;
+ if (collationId == UTF8_BINARY_COLLATION_ID) {
+ comparator = UTF8String::binaryCompare;
+ } else {
+ comparator = (s1, s2) -> {
+ UTF8String convertedS1 = caseAndTrimmingConversionUTF8Binary(s1);
+ UTF8String convertedS2 = caseAndTrimmingConversionUTF8Binary(s2);
+ return convertedS1.binaryCompare(convertedS2);
+ };
+ }
+ return new Collation(
+ collationName(),
+ null,
+ comparator,
+ "1.0",
+ s -> (long) caseAndTrimmingConversionUTF8Binary(s).hashCode(),
+ collationId == UTF8_BINARY_COLLATION_ID,
+ collationId == UTF8_BINARY_COLLATION_ID,
+ collationId == UTF8_BINARY_LCASE_COLLATION_ID
+ );
+ }
+
+ private UTF8String caseAndTrimmingConversionUTF8Binary(UTF8String s) {
Review Comment:
My suggestion would be to only cover collation id format, as you did above.
No need to implement support/parsing for anything but case and accent at this
point.
We can follow up with trimming and co support.
--
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]
