From 5149deddbf1f6fd985fed3c4c2855027a6393c96 Mon Sep 17 00:00:00 2001 From: Menelaos Karavelas <6397014+mkaravel@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:29:32 +0000 Subject: [PATCH] [SQL] Make unary trim collation-aware for case-insensitive ICU collations --- .../sql/catalyst/util/CollationSupport.java | 57 ++++++++++ .../unsafe/types/CollationSupportSuite.java | 100 +++++++++++++++++- .../expressions/stringExpressions.scala | 12 +-- .../expressions/StringExpressionsSuite.scala | 17 +++ 4 files changed, 175 insertions(+), 11 deletions(-) diff --git a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationSupport.java b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationSupport.java index f3f3fe3449a91..43176c3741810 100644 --- a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationSupport.java +++ b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationSupport.java @@ -30,6 +30,12 @@ */ public final class CollationSupport { + private static final UTF8String DEFAULT_TRIM_STRING = UTF8String.fromString(" "); + + private static boolean useCollationAwareDefaultTrim(final int collationId) { + return CollationFactory.isCaseInsensitive(collationId); + } + /** * Collation-aware string expressions. */ @@ -506,6 +512,11 @@ public static class StringTrim { public static UTF8String exec(final UTF8String srcString) { return execBinary(srcString); } + public static UTF8String exec(final UTF8String srcString, final int collationId) { + return useCollationAwareDefaultTrim(collationId) ? + execICU(srcString, collationId) : + execBinary(srcString); + } public static UTF8String exec( final UTF8String srcString, final UTF8String trimString, @@ -527,6 +538,13 @@ public static UTF8String exec( public static String genCode(final String srcString) { return String.format("CollationSupport.StringTrim.execBinary(%s)", srcString); } + public static String genCode(final String srcString, final int collationId) { + if (useCollationAwareDefaultTrim(collationId)) { + return String.format( + "CollationSupport.StringTrim.execICU(%s, %d)", srcString, collationId); + } + return genCode(srcString); + } public static String genCode( final String srcString, final String trimString, @@ -553,6 +571,11 @@ public static UTF8String execLowercase( final int collationId) { return CollationAwareUTF8String.lowercaseTrim(srcString, trimString, collationId); } + public static UTF8String execICU( + final UTF8String srcString, + final int collationId) { + return execICU(srcString, DEFAULT_TRIM_STRING, collationId); + } public static UTF8String execICU( final UTF8String srcString, final UTF8String trimString, @@ -571,6 +594,11 @@ public static class StringTrimLeft { public static UTF8String exec(final UTF8String srcString) { return execBinary(srcString); } + public static UTF8String exec(final UTF8String srcString, final int collationId) { + return useCollationAwareDefaultTrim(collationId) ? + execICU(srcString, collationId) : + execBinary(srcString); + } public static UTF8String exec( final UTF8String srcString, UTF8String trimString, @@ -589,6 +617,13 @@ public static UTF8String exec( public static String genCode(final String srcString) { return String.format("CollationSupport.StringTrimLeft.execBinary(%s)", srcString); } + public static String genCode(final String srcString, final int collationId) { + if (useCollationAwareDefaultTrim(collationId)) { + return String.format( + "CollationSupport.StringTrimLeft.execICU(%s, %d)", srcString, collationId); + } + return genCode(srcString); + } public static String genCode( final String srcString, final String trimString, @@ -613,6 +648,11 @@ public static UTF8String execLowercase( final UTF8String trimString) { return CollationAwareUTF8String.lowercaseTrimLeft(srcString, trimString); } + public static UTF8String execICU( + final UTF8String srcString, + final int collationId) { + return execICU(srcString, DEFAULT_TRIM_STRING, collationId); + } public static UTF8String execICU( final UTF8String srcString, final UTF8String trimString, @@ -625,6 +665,11 @@ public static class StringTrimRight { public static UTF8String exec(final UTF8String srcString) { return execBinary(srcString); } + public static UTF8String exec(final UTF8String srcString, final int collationId) { + return useCollationAwareDefaultTrim(collationId) ? + execICU(srcString, collationId) : + execBinary(srcString); + } public static UTF8String exec( final UTF8String srcString, final UTF8String trimString, @@ -646,6 +691,13 @@ public static UTF8String exec( public static String genCode(final String srcString) { return String.format("CollationSupport.StringTrimRight.execBinary(%s)", srcString); } + public static String genCode(final String srcString, final int collationId) { + if (useCollationAwareDefaultTrim(collationId)) { + return String.format( + "CollationSupport.StringTrimRight.execICU(%s, %d)", srcString, collationId); + } + return genCode(srcString); + } public static String genCode( final String srcString, final String trimString, @@ -671,6 +723,11 @@ public static UTF8String execLowercase( final int collationId) { return CollationAwareUTF8String.lowercaseTrimRight(srcString, trimString, collationId); } + public static UTF8String execICU( + final UTF8String srcString, + final int collationId) { + return execICU(srcString, DEFAULT_TRIM_STRING, collationId); + } public static UTF8String execICU( final UTF8String srcString, final UTF8String trimString, diff --git a/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java b/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java index 7f3c1a925e026..f6d8514a8e653 100644 --- a/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java +++ b/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java @@ -2757,11 +2757,11 @@ private void assertStringTrim(String collationName, String sourceString, String if (trimString == null) { // Trim string is ASCII space. - result = CollationSupport.StringTrim.exec(src); - UTF8String trimLeft = CollationSupport.StringTrimLeft.exec(src); - resultTrimLeftRight = CollationSupport.StringTrimRight.exec(trimLeft); - UTF8String trimRight = CollationSupport.StringTrimRight.exec(src); - resultTrimRightLeft = CollationSupport.StringTrimLeft.exec(trimRight); + result = CollationSupport.StringTrim.exec(src, collationId); + UTF8String trimLeft = CollationSupport.StringTrimLeft.exec(src, collationId); + resultTrimLeftRight = CollationSupport.StringTrimRight.exec(trimLeft, collationId); + UTF8String trimRight = CollationSupport.StringTrimRight.exec(src, collationId); + resultTrimRightLeft = CollationSupport.StringTrimLeft.exec(trimRight, collationId); } else { // Trim string is specified. result = CollationSupport.StringTrim.exec(src, trim, collationId); @@ -2778,6 +2778,96 @@ private void assertStringTrim(String collationName, String sourceString, String assertEquals(resultTrimRightLeft, result); } + private void assertDefaultStringTrims( + String collationName, + String sourceString, + String expectedLeft, + String expectedRight, + String expectedBoth) throws SparkException { + int collationId = CollationFactory.collationNameToId(collationName); + UTF8String source = UTF8String.fromString(sourceString); + UTF8String defaultTrimString = UTF8String.fromString(" "); + + UTF8String trimLeft = CollationSupport.StringTrimLeft.exec(source, collationId); + UTF8String trimRight = CollationSupport.StringTrimRight.exec(source, collationId); + UTF8String trimBoth = CollationSupport.StringTrim.exec(source, collationId); + + assertEquals(UTF8String.fromString(expectedLeft), trimLeft); + assertEquals(UTF8String.fromString(expectedRight), trimRight); + assertEquals(UTF8String.fromString(expectedBoth), trimBoth); + assertEquals( + CollationSupport.StringTrimLeft.exec(source, defaultTrimString, collationId), trimLeft); + assertEquals( + CollationSupport.StringTrimRight.exec(source, defaultTrimString, collationId), trimRight); + assertEquals( + CollationSupport.StringTrim.exec(source, defaultTrimString, collationId), trimBoth); + } + + @Test + public void testDefaultStringTrimsUseCollation() throws SparkException { + String[] spaceSeparators = { + "\u00A0", "\u1680", "\u2000", "\u2001", "\u2002", "\u2003", "\u2004", "\u2005", + "\u2006", "\u2007", "\u2008", "\u2009", "\u200A", "\u202F", "\u205F", "\u3000" + }; + String[] unaffectedCollations = { + UTF8_BINARY, + "UTF8_BINARY_RTRIM", + UTF8_LCASE, + "UTF8_LCASE_RTRIM", + UNICODE, + "UNICODE_RTRIM" + }; + String[] affectedCollations = { + UNICODE_CI, + "UNICODE_CI_RTRIM", + "UNICODE_CI_AI", + "UNICODE_CI_AI_RTRIM" + }; + + for (String collation : unaffectedCollations) { + int collationId = CollationFactory.collationNameToId(collation); + assertEquals( + "CollationSupport.StringTrim.execBinary(source)", + CollationSupport.StringTrim.genCode("source", collationId)); + assertEquals( + "CollationSupport.StringTrimLeft.execBinary(source)", + CollationSupport.StringTrimLeft.genCode("source", collationId)); + assertEquals( + "CollationSupport.StringTrimRight.execBinary(source)", + CollationSupport.StringTrimRight.genCode("source", collationId)); + } + for (String collation : affectedCollations) { + int collationId = CollationFactory.collationNameToId(collation); + assertEquals( + String.format("CollationSupport.StringTrim.execICU(source, %d)", collationId), + CollationSupport.StringTrim.genCode("source", collationId)); + assertEquals( + String.format("CollationSupport.StringTrimLeft.execICU(source, %d)", collationId), + CollationSupport.StringTrimLeft.genCode("source", collationId)); + assertEquals( + String.format("CollationSupport.StringTrimRight.execICU(source, %d)", collationId), + CollationSupport.StringTrimRight.genCode("source", collationId)); + } + + for (String separator : spaceSeparators) { + String source = separator + "abc" + separator; + for (String collation : unaffectedCollations) { + assertDefaultStringTrims(collation, source, source, source, source); + } + for (String collation : affectedCollations) { + assertDefaultStringTrims( + collation, source, "abc" + separator, separator + "abc", "abc"); + } + } + + for (String collation : affectedCollations) { + assertDefaultStringTrims(collation, " abc ", "abc ", " abc", "abc"); + assertDefaultStringTrims(collation, "\tabc\t", "\tabc\t", "\tabc\t", "\tabc\t"); + assertDefaultStringTrims( + collation, "\u200Babc\u200B", "\u200Babc\u200B", "\u200Babc\u200B", "\u200Babc\u200B"); + } + } + @Test public void testStringTrim() throws SparkException { // Basic tests. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala index b918591a43acf..d05c3ab652153 100755 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala @@ -1490,11 +1490,11 @@ trait String2TrimExpression extends Expression with ImplicitCastInputTypes { if (evals.length == 1) { val stringTrimCode: String = this match { case _: StringTrim => - CollationSupport.StringTrim.genCode(srcString.value) + CollationSupport.StringTrim.genCode(srcString.value, collationId) case _: StringTrimLeft => - CollationSupport.StringTrimLeft.genCode(srcString.value) + CollationSupport.StringTrimLeft.genCode(srcString.value, collationId) case _: StringTrimRight => - CollationSupport.StringTrimRight.genCode(srcString.value) + CollationSupport.StringTrimRight.genCode(srcString.value, collationId) } ev.copy(code = code""" |${srcString.code} @@ -1622,7 +1622,7 @@ case class StringTrim(srcStr: Expression, trimStr: Option[Expression] = None) override protected def direction: String = "BOTH" override def doEval(srcString: UTF8String): UTF8String = - CollationSupport.StringTrim.exec(srcString) + CollationSupport.StringTrim.exec(srcString, collationId) override def doEval(srcString: UTF8String, trimString: UTF8String): UTF8String = CollationSupport.StringTrim.exec(srcString, trimString, collationId) @@ -1739,7 +1739,7 @@ case class StringTrimLeft(srcStr: Expression, trimStr: Option[Expression] = None override protected def direction: String = "LEADING" override def doEval(srcString: UTF8String): UTF8String = - CollationSupport.StringTrimLeft.exec(srcString) + CollationSupport.StringTrimLeft.exec(srcString, collationId) override def doEval(srcString: UTF8String, trimString: UTF8String): UTF8String = CollationSupport.StringTrimLeft.exec(srcString, trimString, collationId) @@ -1807,7 +1807,7 @@ case class StringTrimRight(srcStr: Expression, trimStr: Option[Expression] = Non override protected def direction: String = "TRAILING" override def doEval(srcString: UTF8String): UTF8String = - CollationSupport.StringTrimRight.exec(srcString) + CollationSupport.StringTrimRight.exec(srcString, collationId) override def doEval(srcString: UTF8String, trimString: UTF8String): UTF8String = CollationSupport.StringTrimRight.exec(srcString, trimString, collationId) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala index 4ff0a9788d99b..22dc68ab1d8e3 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions.Cast._ import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection import org.apache.spark.sql.catalyst.util.CharsetProvider +import org.apache.spark.sql.catalyst.util.CollationFactory import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLId import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.types.StringTypeWithCollation @@ -1060,6 +1061,22 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { checkEvaluation(StringTrim(Literal("xxxbarxxx"), Literal("x")), "bar") } + test("default trim uses collation-aware space matching") { + val unicodeCi = CollationFactory.collationNameToId("UNICODE_CI") + val unicode = CollationFactory.collationNameToId("UNICODE") + val source = "\u00A0abc\u00A0" + val unicodeCiSource = Literal.create(source, StringType(unicodeCi)) + val unicodeSource = Literal.create(source, StringType(unicode)) + + checkEvaluation(StringTrimLeft(unicodeCiSource), "abc\u00A0") + checkEvaluation(StringTrimRight(unicodeCiSource), "\u00A0abc") + checkEvaluation(StringTrim(unicodeCiSource), "abc") + + checkEvaluation(StringTrimLeft(unicodeSource), source) + checkEvaluation(StringTrimRight(unicodeSource), source) + checkEvaluation(StringTrim(unicodeSource), source) + } + test("LTRIM") { val s = $"a".string.at(0) checkEvaluation(StringTrimLeft(Literal(" aa ")), "aa ", create_row(" abdef "))