You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "uros-db (via GitHub)" <gi...@apache.org> on 2024/05/23 18:17:38 UTC

[PR] [WIP][SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

uros-db opened a new pull request, #46722:
URL: https://github.com/apache/spark/pull/46722

   <!--
   Thanks for sending a pull request!  Here are some tips for you:
     1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
     2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
     3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
     4. Be sure to keep the PR description updated to reflect all changes.
     5. Please write your PR title to summarize what this PR proposes.
     6. If possible, provide a concise example to reproduce the issue for a faster review.
     7. If you want to add a new configuration, please read the guideline first for naming configurations in
        'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
     8. If you want to add or modify an error type or message, please read the guideline first in
        'common/utils/src/main/resources/error/README.md'.
   -->
   
   ### What changes were proposed in this pull request?
   Enable collation support for hash join on complex types.
   
   - Logical plan is rewritten in analysis to (recursively) replace all non-binary strings with CollationKey
   - CollationKey is a unary expression that transforms StringType to BinaryType
   - Collation keys allow correct & efficient string comparison under specific collation rules
   
   
   ### Why are the changes needed?
   Improve JOIN performance for complex types containing collated strings.
   
   
   ### Does this PR introduce _any_ user-facing change?
   No.
   
   
   ### How was this patch tested?
   
   - Unit tests for `CollationKey` in `CollationExpressionSuite`
   - E2e SQL tests for `RewriteCollationJoin` in `CollationSuite`
   - Various queries with JOIN in existing TPCDS collation test suite
   
   
   ### Was this patch authored or co-authored using generative AI tooling?
   No.
   


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "dbatomic (via GitHub)" <gi...@apache.org>.
dbatomic commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1622198569


##########
sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala:
##########
@@ -1046,6 +1046,211 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
     })
   }
 
+  test("hash join should be used for arrays of collated strings") {
+    val t1 = "T_1"
+    val t2 = "T_2"
+
+    case class HashJoinTestCase[R](collation: String, result: R)
+    val testCases = Seq(
+      HashJoinTestCase("UTF8_BINARY",
+        Seq(Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UTF8_BINARY_LCASE",
+        Seq(Row(Seq("aa"), 1, Seq("AA"), 2), Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UNICODE",
+        Seq(Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UNICODE_CI",
+        Seq(Row(Seq("aa"), 1, Seq("AA"), 2), Row(Seq("aa"), 1, Seq("aa"), 2)))
+    )
+
+    testCases.foreach(t => {
+      withTable(t1, t2) {
+        sql(s"CREATE TABLE $t1 (x ARRAY<STRING COLLATE ${t.collation}>, i int) USING PARQUET")
+        sql(s"INSERT INTO $t1 VALUES (array('aa'), 1)")
+
+        sql(s"CREATE TABLE $t2 (y ARRAY<STRING COLLATE ${t.collation}>, j int) USING PARQUET")
+        sql(s"INSERT INTO $t2 VALUES (array('AA'), 2), (array('aa'), 2)")
+
+        val df = sql(s"SELECT * FROM $t1 JOIN $t2 ON $t1.x = $t2.y")
+        checkAnswer(df, t.result)
+
+        val queryPlan = df.queryExecution.executedPlan
+
+        // confirm that hash join is used instead of sort merge join
+        assert(
+          collectFirst(queryPlan) {
+            case _: HashJoin => ()
+          }.nonEmpty
+        )
+        assert(
+          collectFirst(queryPlan) {
+            case _: ShuffledJoin => ()
+          }.isEmpty
+        )
+
+        // if collation doesn't support binary equality, collation key should be injected
+        if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {

Review Comment:
   Can you:
   1) Verify that there is no array transform if collation is not binary.
   2) Check that there is CollationKey as part of ArrayTransform function.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1619442328


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,68 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      case st: StringType
+        if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
+          CollationKey(expr)
+
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+

Review Comment:
   how about map type? I think we support map type join/grouping keys now?



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "uros-db (via GitHub)" <gi...@apache.org>.
uros-db commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1620601656


##########
sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala:
##########
@@ -1046,6 +1046,211 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
     })
   }
 
+  test("hash join should be used for arrays of collated strings") {
+    val t1 = "T_1"
+    val t2 = "T_2"
+
+    case class HashJoinTestCase[R](collation: String, result: R)
+    val testCases = Seq(
+      HashJoinTestCase("UTF8_BINARY",
+        Seq(Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UTF8_BINARY_LCASE",
+        Seq(Row(Seq("aa"), 1, Seq("AA"), 2), Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UNICODE",
+        Seq(Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UNICODE_CI",

Review Comment:
   I agree that we should start adding these tests. However, since this PR is kind of an extension on the previous one (https://github.com/apache/spark/pull/46599), I would say it may be better to proceed now, and then add tests for both later? we could possibly dedicate one PR to improve testing for CollationSuite with respect to new collations - let me know if that sounds good



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "dbatomic (via GitHub)" <gi...@apache.org>.
dbatomic commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1620584813


##########
sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala:
##########
@@ -1046,6 +1046,211 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
     })
   }
 
+  test("hash join should be used for arrays of collated strings") {
+    val t1 = "T_1"
+    val t2 = "T_2"
+
+    case class HashJoinTestCase[R](collation: String, result: R)
+    val testCases = Seq(
+      HashJoinTestCase("UTF8_BINARY",
+        Seq(Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UTF8_BINARY_LCASE",
+        Seq(Row(Seq("aa"), 1, Seq("AA"), 2), Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UNICODE",
+        Seq(Row(Seq("aa"), 1, Seq("aa"), 2))),
+      HashJoinTestCase("UNICODE_CI",

Review Comment:
   Can we start adding at least one locale + case insensitive + accent insensitive variation in tests? E.g. I would like to see if things work correctly for SRB_CI_AI, SRB_CS_AS, SRB_CI_AS and SRB_CS_AI.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "uros-db (via GitHub)" <gi...@apache.org>.
uros-db commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1622544544


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,68 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      case st: StringType
+        if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
+          CollationKey(expr)
+
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+
+      case _ =>
+        expr
+    }
+  }
+
+  private def processStruct(str: Expression, fields: Array[StructField]): Expression = {
+    val struct = CreateNamedStruct(fields.zipWithIndex.flatMap {

Review Comment:
   I think `case` is valid here, I'm not aware of another syntax that handles this properly



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1624959243


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,83 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.catalyst.util.UnsafeRowUtils
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
+/**
+ * This rule rewrites Join conditions to ensure that all types containing non-binary collated
+ * strings are compared correctly. This is necessary because join conditions are evaluated using
+ * binary equality, which does not work correctly for non-binary collated strings. However, by
+ * injecting CollationKey expressions into the join condition, we can ensure that the comparison
+ * is done correctly, which then allows HashJoin to work properly on this type of data.
+ */
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  /**
+   * Recursively process the expression in order to replace non-binary collated strings with their
+   * associated collation keys. This is necessary to ensure that the join condition is evaluated
+   * correctly for all types containing non-binary collated strings, including structs and arrays.
+   */
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      // For binary stable expressions, no special handling is needed.
+      case _ if UnsafeRowUtils.isBinaryStable(dt) =>
+        expr
+
+      // Inject CollationKey for non-binary collated strings.
+      case _: StringType =>
+          CollationKey(expr)
+
+      // Recursively process struct fields for non-binary structs.
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      // Recursively process array elements for non-binary arrays.
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+
+      // Joins are not supported on maps, so there's no special handling for MapType.
+      case _ =>
+        expr
+    }
+  }
+
+  private def processStruct(str: Expression, fields: Array[StructField]): Expression = {
+    val struct = CreateNamedStruct(fields.zipWithIndex.flatMap { case (f, i) =>
+      Seq(Literal(f.name), processExpression(GetStructField(str, i, Some(f.name)), f.dataType))
+    }.toImmutableArraySeq)

Review Comment:
   shall we return the original `str` expression if no field is updated with CollationKey?



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan closed pull request #46722: [SPARK-48318][SQL] Enable hash join support for all collations (complex types)
URL: https://github.com/apache/spark/pull/46722


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "dbatomic (via GitHub)" <gi...@apache.org>.
dbatomic commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1622351854


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,68 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      case st: StringType
+        if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
+          CollationKey(expr)
+
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+
+      case _ =>
+        expr
+    }
+  }
+
+  private def processStruct(str: Expression, fields: Array[StructField]): Expression = {
+    val struct = CreateNamedStruct(fields.zipWithIndex.flatMap {
+      case (f, i) =>
+        Seq(Literal(f.name),
+          processExpression(GetStructField(str, i, Some(f.name)), f.dataType))
+    }.toImmutableArraySeq)
+    if (struct.valExprs.forall(_.isInstanceOf[GetStructField])) {

Review Comment:
   So, you are checking here whether all the expressions remained `GetStructField`, right?
   Maybe it is cleaner to call `isBinaryStable` against dataType first. And if false call `processExpression`?
   
   Also, `processExpression` should be called something more descriptive (`insertCollationKey` or something like that).



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "uros-db (via GitHub)" <gi...@apache.org>.
uros-db commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1620319639


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,68 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      case st: StringType
+        if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
+          CollationKey(expr)
+
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+

Review Comment:
   I don't think joins and aggregations are supported for map type at this time
   
   join:
   ```
   Failed to analyze query: org.apache.spark.sql.AnalysisException: [DATATYPE_MISMATCH.INVALID_ORDERING_TYPE] Cannot resolve "(m = m)" due to data type mismatch: The `=` does not support ordering on type "MAP<STRING, STRING>"
   ```
   
   agg:
   ```
   Failed to analyze query: org.apache.spark.sql.AnalysisException: [UNSUPPORTED_FEATURE.SET_OPERATION_ON_MAP_TYPE] The feature is not supported: Cannot have MAP type columns in DataFrame which calls set operations (INTERSECT, EXCEPT, etc.), but the type of column `m` is "MAP<STRING, STRING>".
   ```



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "uros-db (via GitHub)" <gi...@apache.org>.
uros-db commented on PR #46722:
URL: https://github.com/apache/spark/pull/46722#issuecomment-2142732488

   @cloud-fan ready for further review


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "dbatomic (via GitHub)" <gi...@apache.org>.
dbatomic commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1622346805


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,68 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      case st: StringType
+        if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
+          CollationKey(expr)
+
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+
+      case _ =>
+        expr
+    }
+  }
+
+  private def processStruct(str: Expression, fields: Array[StructField]): Expression = {
+    val struct = CreateNamedStruct(fields.zipWithIndex.flatMap {

Review Comment:
   do you need case here? Why not just flatMap((f,i) => { })?



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "uros-db (via GitHub)" <gi...@apache.org>.
uros-db commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1622570360


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,68 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      case st: StringType
+        if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
+          CollationKey(expr)
+
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+
+      case _ =>
+        expr
+    }
+  }
+
+  private def processStruct(str: Expression, fields: Array[StructField]): Expression = {
+    val struct = CreateNamedStruct(fields.zipWithIndex.flatMap {
+      case (f, i) =>
+        Seq(Literal(f.name),
+          processExpression(GetStructField(str, i, Some(f.name)), f.dataType))
+    }.toImmutableArraySeq)
+    if (struct.valExprs.forall(_.isInstanceOf[GetStructField])) {

Review Comment:
   Or even better, we can short-circuit right away for all binary stable expressions in `processExpression`.
   
   As for the naming, all these methods are private and internal to `RewriteCollationJoin` - I've added comments explaining what's going on, so I hope this should be enough



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "dbatomic (via GitHub)" <gi...@apache.org>.
dbatomic commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1622190538


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,68 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  private def processExpression(expr: Expression, dt: DataType): Expression = {

Review Comment:
   Hey, can you add a comment explaining what is happening here.
   I think that it is not obvious why we are inserting `CollationKey` as a rule. I would appreciate comment header on top of `RewriteCollationJoin` that explains why and what we are doing here.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "uros-db (via GitHub)" <gi...@apache.org>.
uros-db commented on code in PR #46722:
URL: https://github.com/apache/spark/pull/46722#discussion_r1625384664


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteCollationJoin.scala:
##########
@@ -17,29 +17,83 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationKey, Equality}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.catalyst.util.UnsafeRowUtils
+import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
 
+/**
+ * This rule rewrites Join conditions to ensure that all types containing non-binary collated
+ * strings are compared correctly. This is necessary because join conditions are evaluated using
+ * binary equality, which does not work correctly for non-binary collated strings. However, by
+ * injecting CollationKey expressions into the join condition, we can ensure that the comparison
+ * is done correctly, which then allows HashJoin to work properly on this type of data.
+ */
 object RewriteCollationJoin extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case j @ Join(_, _, _, Some(condition), _) =>
       val newCondition = condition transform {
         case e @ Equality(l: AttributeReference, r: AttributeReference) =>
-          (l.dataType, r.dataType) match {
-            case (st: StringType, _: StringType)
-              if !CollationFactory.fetchCollation(st.collationId).supportsBinaryEquality =>
-                e.withNewChildren(Seq(CollationKey(l), CollationKey(r)))
-            case _ =>
-              e
-          }
+          e.withNewChildren(Seq(processExpression(l, l.dataType), processExpression(r, r.dataType)))
       }
       if (!newCondition.fastEquals(condition)) {
         j.copy(condition = Some(newCondition))
       } else {
         j
       }
   }
+
+  /**
+   * Recursively process the expression in order to replace non-binary collated strings with their
+   * associated collation keys. This is necessary to ensure that the join condition is evaluated
+   * correctly for all types containing non-binary collated strings, including structs and arrays.
+   */
+  private def processExpression(expr: Expression, dt: DataType): Expression = {
+    dt match {
+      // For binary stable expressions, no special handling is needed.
+      case _ if UnsafeRowUtils.isBinaryStable(dt) =>
+        expr
+
+      // Inject CollationKey for non-binary collated strings.
+      case _: StringType =>
+          CollationKey(expr)
+
+      // Recursively process struct fields for non-binary structs.
+      case StructType(fields) =>
+        processStruct(expr, fields)
+
+      // Recursively process array elements for non-binary arrays.
+      case ArrayType(et, containsNull) =>
+        processArray(expr, et, containsNull)
+
+      // Joins are not supported on maps, so there's no special handling for MapType.
+      case _ =>
+        expr
+    }
+  }
+
+  private def processStruct(str: Expression, fields: Array[StructField]): Expression = {
+    val struct = CreateNamedStruct(fields.zipWithIndex.flatMap { case (f, i) =>
+      Seq(Literal(f.name), processExpression(GetStructField(str, i, Some(f.name)), f.dataType))
+    }.toImmutableArraySeq)

Review Comment:
   `processExpression` should handle that by short-circuiting in the first `case`
   
   furthermore, if all fields in a struct will remain unchanged - `processStruct` won't even get called to begin with



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-48318][SQL] Enable hash join support for all collations (complex types) [spark]

Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #46722:
URL: https://github.com/apache/spark/pull/46722#issuecomment-2146772376

   thanks, merging to master!


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org