[SPARK-47541][SQL] Collated strings in complex types supporting operations reverse, array_join, concat, map by nikolamand-db · Pull Request #45693 · apache/spark

nikolamand-db · 2024-03-25T09:04:34Z

What changes were proposed in this pull request?

Add proper support for complex types containing collated strings in operations reverse, array_join, concat, map (create). Examples:

select reverse('abc' collate utf8_binary_lcase);
select reverse(array('a' collate utf8_binary_lcase, 'b' collate utf8_binary_lcase));
select array_join(array('a' collate utf8_binary_lcase, 'b' collate utf8_binary_lcase), ', ' collate utf8_binary_lcase);
select concat('a' collate utf8_binary_lcase, 'b' collate utf8_binary_lcase);
select map('a' collate utf8_binary_lcase, 1, 'A' collate utf8_binary_lcase, 2);

Why are the changes needed?

To enable listed complex types operations support for collated strings.

Does this PR introduce any user-facing change?

Yes, results of listed complex types operations will return expected results to users.

How was this patch tested?

Added checks to collations suite.

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

No.

mihailomilosevic2001 · 2024-03-25T09:21:10Z

    array: Expression,
    delimiter: Expression,
-    nullReplacement: Option[Expression]) extends Expression with ExpectsInputTypes {
+    nullReplacement: Option[Expression]) extends ImplicitCastInputTypes with ExpectsInputTypes {


You extend both Parent class and its Child. Remove ExpectsInputTypes.

mihailomilosevic2001 · 2024-03-25T09:25:09Z

+  private val commonType: DataType = {
+    val elementType = array.dataType.asInstanceOf[ArrayType].elementType
+    val s = nullReplacement match {
+      case Some(replacement) => Seq(elementType, delimiter.dataType, replacement.dataType)
+      case _ => Seq(elementType, delimiter.dataType)
+    }
+    TypeCoercion.findWiderCommonType(s).getOrElse(StringType)
+  }
+
  override def inputTypes: Seq[AbstractDataType] = if (nullReplacement.isDefined) {
-    Seq(ArrayType(StringType), StringType, StringType)
+    Seq(ArrayType(commonType), commonType, commonType)
  } else {
-    Seq(ArrayType(StringType), StringType)
+    Seq(ArrayType(commonType), commonType)
  }


This seems a bit too much hustle to get expected input types. I would say these have to not depend on the parameters. Why not use StringTypeAnyCollation?

One more thing is. This what you do with findWiderCommonType is not correct. If delimiter has explicit collation, we would like to have that collation. In your case you would collect two different collations, where arrayType.elementType would go first and findWiderCommonType practically does not do anything with delimiter's type, but would keep the first element type.

mihailomilosevic2001 · 2024-03-25T09:26:39Z

  }

-  override def dataType: DataType = StringType
+  override def dataType: DataType = commonType


Why not use some type of pass through where we would say dataType = array.asInstanceOf[ArrayType].elementType? Implicit casting should solve the problem of collations.

mihailomilosevic2001 · 2024-03-25T11:27:36Z

+  override def inputTypes: Seq[AbstractDataType] = {
+    val arrayType = array.dataType.asInstanceOf[ArrayType].elementType match {
+      case _: StringType => array.dataType
+      case _ => ArrayType(StringType)
+    }
+    if (nullReplacement.isDefined) {
+      Seq(arrayType, StringTypeAnyCollation, StringTypeAnyCollation)
+    } else {
+      Seq(arrayType, StringTypeAnyCollation)
+    }


This looks better, but I am still concerned that we are setting inputTypes according to inputs that we get. If I understand the logic correctly, it should be independent. I would rather leave this as ArrayType only and then override checkInputTypes to do this check you performed.

@mihailom-db There are already examples of such checks performed in inputTypes, here are few:

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala

Lines 59 to 68 in 49b7dbc

override def inputTypes: Seq[AbstractDataType] = {

(left.dataType, right.dataType) match {

case (ArrayType(e1, hasNull1), ArrayType(e2, hasNull2)) =>

TypeCoercion.findTightestCommonType(e1, e2) match {

case Some(dt) => Seq(ArrayType(dt, hasNull1), ArrayType(dt, hasNull2))

case _ => Seq.empty

}

case _ => Seq.empty

}

}

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala

Lines 2341 to 2350 in 49b7dbc

override def inputTypes: Seq[AbstractDataType] = {

(left.dataType, right.dataType) match {

case (ArrayType(e1, hasNull), e2) =>

TypeCoercion.findTightestCommonType(e1, e2) match {

case Some(dt) => Seq(ArrayType(dt, hasNull), dt)

case _ => Seq.empty

}

case _ => Seq.empty

}

}

I see a potential issue in implicit conversions if I replace it with ArrayType without StringType information. Is there some kind of explicit conversion which will override this?

ArrayJoin should be fine with only ArrayType. This is because TypeCoercion rules in FunctionArgumentConversion cast this argument to ArrayType(StringType) and catalyst rules will loop over this and cast it to proper collation if necessary.

dbatomic · 2024-03-26T11:08:26Z

    // Binary type data is `byte[]`, which can't use `==` to check equality.
    case _: AtomicType | _: CalendarIntervalType | _: NullType
-      if !keyType.isInstanceOf[BinaryType] => new java.util.HashMap[Any, Int]()
+      if !keyType.isInstanceOf[BinaryType] && (!keyType.isInstanceOf[StringType] ||


Let's wait for Collation Refactory PR for this. This should be stringType.isBinaryComparable?

Also, can you push StringType in a separate case. Relying on order evaluation of || is not a good thing (I don't know if scala standard enforces this). E.g. you can end up with keyType.asInstanceOf[StringType].isBinaryCollation evaluating first which will throw if this is not a StringType.

mihailomilosevic2001

Seems good to me in combination with #45383

cloud-fan · 2024-04-04T09:39:38Z


  override def inputTypes: Seq[AbstractDataType] = if (nullReplacement.isDefined) {
-    Seq(ArrayType(StringType), StringType, StringType)
+    Seq(ArrayType, StringTypeAnyCollation, StringTypeAnyCollation)


what if the input is array of int? Before it will fail with type check, but now it seems different.

We probably need to override checkInputDataTypes to check the array element.

@cloud-fan There's a special case in FunctionArgumentConversion which implicitly casts array parameter to array of strings #21620.

So if I'm not mistaken this behavior was allowed before and it still works as expected.

cloud-fan

LGTM except for one issue

cloud-fan · 2024-04-08T03:04:02Z

thanks, merging to master!

### What changes were proposed in this pull request? Addition of abstract arraytype which accepts StringTypeCollated as elementType. Changes in this PR #45693 work for ArrayJoin, but will not work in general for other functions. This PR introduces a change to give an interface for all functions. Merge only after #45693. ### Why are the changes needed? This is needed in order to enable functions to use collated arrays. ### Does this PR introduce _any_ user-facing change? Yes, collation functions will work. ### How was this patch tested? Test for array_join added to `CollationSuite` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #45891 from mihailom-db/SPARK-47736. Authored-by: Mihailo Milosevic <mihailo.milosevic@databricks.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com>

dongjoon-hyun

Hi, @nikolamand-db , @dbatomic and @cloud-fan .

Unfortunately, this broke master branch ANSI CI.

https://github.com/apache/spark/actions/workflows/build_ansi.yml

Please check the following and recover the ANSI CI environment.

$ SPARK_ANSI_SQL_MODE=true build/sbt "sql/testOnly *CollationSuite -- -z complex"
...
[info] - Support operations on complex types containing collated strings *** FAILED *** (144 milliseconds)
[info]   java.lang.RuntimeException: Max iterations (100) reached for batch Resolution, please set 'spark.sql.analyzer.maxIterations' to a larger value.

dongjoon-hyun · 2024-04-11T16:55:26Z

Since there are many Collated-feature patch, I made a following PR to catch this kind of failure more eagerly.

[SPARK-47820][INFRA] Run ANSI SQL CI twice per day #46010

Anyway, thank you all for working on this feature.

mihailomilosevic2001 · 2024-04-11T19:35:03Z

I have a fix for this, will create a followup soon.

nikolamand-db added 3 commits March 22, 2024 20:36

Support collations in complex types operations

614c108

Add test

e4c589b

Add checks

b027d8b

github-actions Bot added the SQL label Mar 25, 2024

mihailomilosevic2001 reviewed Mar 25, 2024

View reviewed changes

Cast fixes in array_join

49b7dbc

mihailomilosevic2001 reviewed Mar 25, 2024

View reviewed changes

nikolamand-db added 2 commits March 25, 2024 16:08

Fix error messages and return to expected type in array_join

af5b910

Fixes

8a3b328

dbatomic reviewed Mar 26, 2024

View reviewed changes

nikolamand-db added 2 commits April 2, 2024 12:44

Merge branch 'master' into SPARK-47541

a3c199b

Improve ArrayBasedMapBuilder keyToIndex

7d3a5bd

dbatomic approved these changes Apr 2, 2024

View reviewed changes

Update ArrayJoin inputTypes

231d28f

mihailomilosevic2001 approved these changes Apr 3, 2024

View reviewed changes

cloud-fan reviewed Apr 4, 2024

View reviewed changes

cloud-fan approved these changes Apr 4, 2024

View reviewed changes

mihailomilosevic2001 mentioned this pull request Apr 5, 2024

[SPARK-47736][SQL] Add support for AbstractArrayType #45891

Closed

cloud-fan closed this in f576b85 Apr 8, 2024

dongjoon-hyun reviewed Apr 11, 2024

View reviewed changes

mihailomilosevic2001 mentioned this pull request Apr 11, 2024

[SPARK-47541][SQL][FOLLOWUP] Fix AnsiTypeCoercion to handle ArrayType #46016

Closed

	override def inputTypes: Seq[AbstractDataType] = {
	(left.dataType, right.dataType) match {
	case (ArrayType(e1, hasNull1), ArrayType(e2, hasNull2)) =>
	TypeCoercion.findTightestCommonType(e1, e2) match {
	case Some(dt) => Seq(ArrayType(dt, hasNull1), ArrayType(dt, hasNull2))
	case _ => Seq.empty
	}
	case _ => Seq.empty
	}
	}

	override def inputTypes: Seq[AbstractDataType] = {
	(left.dataType, right.dataType) match {
	case (ArrayType(e1, hasNull), e2) =>
	TypeCoercion.findTightestCommonType(e1, e2) match {
	case Some(dt) => Seq(ArrayType(dt, hasNull), dt)
	case _ => Seq.empty
	}
	case _ => Seq.empty
	}
	}

Conversation

nikolamand-db commented Mar 25, 2024

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?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mihailomilosevic2001 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cloud-fan left a comment

Choose a reason for hiding this comment

Uh oh!

cloud-fan commented Apr 8, 2024

Uh oh!

dongjoon-hyun left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun commented Apr 11, 2024

Uh oh!

mihailomilosevic2001 commented Apr 11, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

dongjoon-hyun left a comment •

edited

Loading