[SPARK-47483][SQL] Add support for aggregation and join operations on arrays of collated strings #45611

nikolamand-db · 2024-03-20T09:58:59Z

What changes were proposed in this pull request?

Example of aggregation sequence:

create table t(a array<string collate utf8_binary_lcase>) using parquet;

insert into t(a) values(array('a' collate utf8_binary_lcase));
insert into t(a) values(array('A' collate utf8_binary_lcase));

select distinct a from t;

Example of join sequence:

create table l(a array<string collate utf8_binary_lcase>) using parquet;
create table r(a array<string collate utf8_binary_lcase>) using parquet;

insert into l(a) values(array('a' collate utf8_binary_lcase));
insert into r(a) values(array('A' collate utf8_binary_lcase));

select * from l join r where l.a = r.a;

Both runs should yield one row since the arrays are considered equal.

Problem is in isBinaryStable function which should return false if any of its subtypes is non-binary collated string.

Why are the changes needed?

To support aggregates and joins in arrays of collated strings properly.

Does this PR introduce any user-facing change?

Yes, it fixes the described scenarios.

How was this patch tested?

Added new checks to collation suite.

Was this patch authored or co-authored using generative AI tooling?

No.

… strings

sql/core/src/test/resources/sql-tests/inputs/collations.sql

dbatomic · 2024-03-20T10:46:49Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

+    val tableName = "test_agg_arr_collated"
+    val simple = Seq(
+      // binary
+      ("utf8_binary", Seq("array('aaa')", "array('AAA')"),


Can you use named case objects instead of tuples? Not sure about the others, but I can't read this :)

dbatomic · 2024-03-20T10:50:36Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

+      ("utf8_binary", Seq("array('aaa', 'bbb')", "array('AAA', 'BBB')"),
+        Seq((Seq("aaa", "bbb"), 1), (Seq("AAA", "BBB"), 1))),
+      ("utf8_binary", Seq("array('aaa')", "array('bbb')", "array('AAA')", "array('BBB')"),
+        Seq((Seq("aaa"), 1), (Seq("bbb"), 1), (Seq("AAA"), 1), (Seq("BBB"), 1))),


Do we want to cover UTF_BINARY test case here at all? In previous examples we mainly used UTF8_BINARY as validation but the real thing we want to test here are collations other than UTF8_BINARY.

dbatomic · 2024-03-20T10:51:33Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

+          }
+        }
+    }
+  }


LGTM, but I would prefer if you could simplify the test cases.

dbatomic · 2024-03-20T10:56:02Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

+            count.map{ case (aggStr, _) => Row(aggStr)})
+        }
+    }
+  }


The test pattern that can be useful here is what we did for Window Aggs.
In short, just create a query that targets mixed-case data with LCASE collation (e.g. "aA", "aa", "AA") and query that targets normalized data with UTF8_BINARY ("aa", "aa", "aa"). Aggs and Joins should return the same result.

You can find test example here:
#45568

Updated tests to support this and added case classes for better readability. Please check again.

dbatomic · 2024-03-20T14:33:40Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

+          // binary collation with values converted to lowercase should match the results as well
+          sql(s"create table $tableNameLowercase(a ${check.dataType}) using parquet")
+          check.rows.map(row =>
+            sql(s"insert into $tableNameLowercase(a) values(${row.toLowerCase(Locale.ROOT)})"))


Can you use UTF8String toLowerCase function instead of java.util.locale? For this case it is ok, but we shouldn't start mixing ICU collator with Java's default one.

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

cloud-fan · 2024-03-21T14:12:33Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

+    ).map((check: AggCheck) =>
+        withTable(tableName, tableNameLowercase) {
+          def checkResults(table: String): Unit = {
+            checkAnswer(sql(s"select a, count(*) from $table group by a"),


It's a bad practice that we try to test all cases in end-to-end tests. We should improve the test coverage in the unit test.

I suggest we add tests in UnsafeRowUtilsSuite to test the isBinaryStable function with different cases: array of string, array of array of string, struct of array of string, etc. The end-to-end test here should just run a few queries to show it works.

MaxGekk · 2024-03-21T16:27:16Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/UnsafeRowUtils.scala

-  def isBinaryStable(dataType: DataType): Boolean = dataType.existsRecursively {
-    case st: StringType => CollationFactory.fetchCollation(st.collationId).isBinaryCollation
-    case _ => true
+  def isBinaryStable(dataType: DataType): Boolean = !dataType.existsRecursively {


why isBinaryStable is in UnsafeRowUtils. Is the implementation bound somehow to unsafe row?

Why it is not in DataTypeUtils or Collation..., for example?

Not sure what are the consequences of moving this function, do you know if we can do that @dbatomic?

MaxGekk · 2024-03-21T16:40:58Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

+        Seq("array(array('AAA'), array('CCC'))"),
+        Seq()
+      )
+    ).map((check: JoinCheck) =>


Better to use .foreach here. map is usually used to build new collection which is not need in this case.

MaxGekk · 2024-03-21T16:43:42Z

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala

@@ -640,6 +641,201 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
        "reason" -> "generation expression cannot contain non-default collated string type"))
  }

+  trait ArrayCheck {


Why don't you put the new code to the end of the test suite?

nikolamand-db · 2024-03-21T19:20:27Z

@MaxGekk @cloud-fan please check the updated tests.

MaxGekk · 2024-03-22T07:14:11Z

+1, LGTM. Merging to master.
Thank you, @nikolamand-db and @dbatomic @cloud-fan @HyukjinKwon for review.

… arrays of collated strings ### What changes were proposed in this pull request? Example of aggregation sequence: ``` create table t(a array<string collate utf8_binary_lcase>) using parquet; insert into t(a) values(array('a' collate utf8_binary_lcase)); insert into t(a) values(array('A' collate utf8_binary_lcase)); select distinct a from t; ``` Example of join sequence: ``` create table l(a array<string collate utf8_binary_lcase>) using parquet; create table r(a array<string collate utf8_binary_lcase>) using parquet; insert into l(a) values(array('a' collate utf8_binary_lcase)); insert into r(a) values(array('A' collate utf8_binary_lcase)); select * from l join r where l.a = r.a; ``` Both runs should yield one row since the arrays are considered equal. Problem is in `isBinaryStable` function which should return false if **any** of its subtypes is non-binary collated string. ### Why are the changes needed? To support aggregates and joins in arrays of collated strings properly. ### Does this PR introduce _any_ user-facing change? Yes, it fixes the described scenarios. ### How was this patch tested? Added new checks to collation suite. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#45611 from nikolamand-db/SPARK-47483. Authored-by: Nikola Mandic <nikola.mandic@databricks.com> Signed-off-by: Max Gekk <max.gekk@gmail.com>

Add support for aggregation and join operations on arrays of collated…

f69b093

… strings

github-actions bot added the SQL label Mar 20, 2024

nikolamand-db changed the title ~~[SPARK-47483] Add support for aggregation and join operations on arrays of collated strings~~ [SPARK-47483][SQL] Add support for aggregation and join operations on arrays of collated strings Mar 20, 2024

Combine withTable into one statement

ca7664e

dbatomic reviewed Mar 20, 2024

View reviewed changes

sql/core/src/test/resources/sql-tests/inputs/collations.sql Outdated Show resolved Hide resolved

Remove checks from golden file

f78557d

dbatomic reviewed Mar 20, 2024

View reviewed changes

Improve tests

824b92a

nikolamand-db requested a review from dbatomic March 20, 2024 14:23

dbatomic reviewed Mar 20, 2024

View reviewed changes

dbatomic approved these changes Mar 20, 2024

View reviewed changes

nikolamand-db added 2 commits March 20, 2024 15:49

Replace java lowercase function with UTF8String

78e8310

Merge branch 'master' into SPARK-47483

588cca1

HyukjinKwon reviewed Mar 21, 2024

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala Outdated Show resolved Hide resolved

HyukjinKwon reviewed Mar 21, 2024

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala Outdated Show resolved Hide resolved

Refactor

5f59648

nikolamand-db requested a review from HyukjinKwon March 21, 2024 11:49

cloud-fan reviewed Mar 21, 2024

View reviewed changes

MaxGekk reviewed Mar 21, 2024

View reviewed changes

MaxGekk mentioned this pull request Mar 21, 2024

[SPARK-47358][SQL][COLLATION] Improve repeat expression support to return correct datatype #45615

Closed

Refactor tests

514d2cd

MaxGekk approved these changes Mar 22, 2024

View reviewed changes

MaxGekk closed this in 8cba15e Mar 22, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPARK-47483][SQL] Add support for aggregation and join operations on arrays of collated strings #45611

[SPARK-47483][SQL] Add support for aggregation and join operations on arrays of collated strings #45611

nikolamand-db commented Mar 20, 2024 •

edited

Loading

dbatomic Mar 20, 2024

dbatomic Mar 20, 2024

dbatomic Mar 20, 2024

dbatomic Mar 20, 2024

nikolamand-db Mar 20, 2024

dbatomic Mar 20, 2024

cloud-fan Mar 21, 2024

MaxGekk Mar 21, 2024 •

edited

Loading

nikolamand-db Mar 21, 2024

MaxGekk Mar 21, 2024

MaxGekk Mar 21, 2024

nikolamand-db commented Mar 21, 2024

MaxGekk commented Mar 22, 2024

[SPARK-47483][SQL] Add support for aggregation and join operations on arrays of collated strings #45611

[SPARK-47483][SQL] Add support for aggregation and join operations on arrays of collated strings #45611

Conversation

nikolamand-db commented Mar 20, 2024 • edited Loading

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

MaxGekk Mar 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nikolamand-db commented Mar 21, 2024

MaxGekk commented Mar 22, 2024

nikolamand-db commented Mar 20, 2024 •

edited

Loading

MaxGekk Mar 21, 2024 •

edited

Loading