viirya commented on code in PR #6170:
URL: https://github.com/apache/datafusion-comet/pull/6170#discussion_r4209880058
##########
native/spark-expr/src/array_funcs/array_extrema.rs:
##########
@@ -53,8 +105,33 @@ impl SparkArrayExtrema {
} else {
array_max_udf()
},
+ string_collations: Vec::new(),
+ unicode_version: 0,
}
}
+
+ /// Scala validates the Unicode version and lists string-leaf collations
depth-first.
+ pub fn with_collations(
Review Comment:
The doc comment says Scala validates the Unicode version, but this
constructor accepts any value. With `UTF8_LCASE` and a version other than 16 or
17, it quietly uses the Unicode 16 table. Could it return an error in that
case, so a mismatch between the serde and the native side fails loudly instead
of returning a different winner than Spark?
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -256,16 +256,57 @@ object CometArrayIntersect
private object ArrayExtremaSupport extends CometTypeShim {
val incompatReason: String =
- "Array extrema use binary string ordering for non-UTF8_BINARY collations "
+
+ "Array extrema support UTF8_BINARY, UTF8_LCASE, and UTF8_LCASE_RTRIM
natively; " +
+ "UTF8_LCASE requires Unicode 16 or 17. Other collations use binary
ordering when opted in " +
"(https://github.com/apache/datafusion-comet/issues/4496)."
- def getSupportLevel(elementType: DataType): SupportLevel = {
- if (hasNonDefaultStringCollation(elementType)) {
- // The dispatcher runs Spark's own comparison with the original
collation IDs, including
- // strings nested in arrays or structs. Keep the native bytewise
comparison opt-in only.
- Incompatible(Some(incompatReason))
- } else {
+ val binaryRtrimReason: String =
+ "UTF8_BINARY_RTRIM array extrema use Spark's comparator."
+
+ private def stringCollations(dt: DataType): Seq[String] = dt match {
+ case stringType: StringType => Seq(stringCollationName(stringType))
+ case ArrayType(elementType, _) => stringCollations(elementType)
+ case StructType(fields) => fields.toSeq.flatMap(f =>
stringCollations(f.dataType))
+ case _ => Seq.empty
+ }
+
+ private def supportsCollations(collations: Seq[String]): Boolean =
collations.forall {
+ case "UTF8_BINARY" | "UTF8_BINARY_RTRIM" => true
+ case "UTF8_LCASE" | "UTF8_LCASE_RTRIM" =>
+ collationUnicodeVersion == 16 || collationUnicodeVersion == 17
+ case _ => false
+ }
+
+ def getSupportLevel(elementType: DataType): SupportLevel =
+ // Keep this path on the dispatcher until its native performance
regression is resolved.
+ // Use Unsupported so allowIncompatible cannot silently change its
comparison semantics.
+ if (stringCollations(elementType).contains("UTF8_BINARY_RTRIM")) {
+ Unsupported(Some(binaryRtrimReason))
+ } else if (supportsCollations(stringCollations(elementType))) {
Compatible()
+ } else {
+ Incompatible(Some(incompatReason))
+ }
+
+ def convert(
+ expr: Expression,
+ inputs: Seq[Attribute],
+ binding: Boolean,
+ isMin: Boolean): Option[ExprOuterClass.Expr] = {
+ val collations = stringCollations(expr.dataType)
+ exprToProtoInternal(expr.children.head, inputs, binding).map { child =>
+ val builder = ExprOuterClass.ArrayExtrema
+ .newBuilder()
+ .setChild(child)
+ .setIsMin(isMin)
+ // Preserve the existing binary comparison when unsupported collations
are explicitly
+ // opted in. By default those expressions use the JVM dispatcher.
+ if (!collations.forall(_ == "UTF8_BINARY") &&
supportsCollations(collations)) {
Review Comment:
When the element type mixes a supported and an unsupported collation, for
example a struct with a `UTF8_LCASE` field and a `UNICODE_CI` field, and the
user opts in, `supportsCollations` is false for the whole list. No collations
are sent, so the `UTF8_LCASE` field is also compared bytewise. The incompatible
reason says only the other collations use binary ordering. Would it work to
send the list with each unsupported entry mapped to `UTF8_BINARY`? That keeps
the reason accurate and keeps the supported leaves correct under opt-in.
##########
native/spark-expr/src/array_funcs/array_extrema.rs:
##########
@@ -293,5 +389,100 @@ fn take_extrema_values(values: &ArrayRef, indices:
&UInt32Array) -> Result<Array
}
}
+/// Build one comparator per child array, adding string collation to Spark's
ordering.
+/// Noncollated values use the shared Spark comparator, including its
floating-point rules.
+fn collated_comparator(
+ array: &ArrayRef,
+ collations: &mut std::slice::Iter<Utf8Collation>,
+ unicode_version: u32,
+) -> Result<DynComparator> {
+ if collations.as_slice().is_empty() {
+ return spark_comparator(array.as_ref(), array.as_ref());
+ }
+ match array.data_type() {
+ DataType::List(_) => {
+ let array = array.as_list::<i32>();
+ let compare = collated_comparator(array.values(), collations,
unicode_version)?;
+ let offsets = array.offsets().clone();
+ Ok(nulls_first(array.nulls().cloned(), move |left, right| {
+ let left_start = offsets[left] as usize;
+ let right_start = offsets[right] as usize;
+ let left_len = offsets[left + 1] as usize - left_start;
+ let right_len = offsets[right + 1] as usize - right_start;
+ for offset in 0..left_len.min(right_len) {
+ let ordering = compare(left_start + offset, right_start +
offset);
+ if ordering != Ordering::Equal {
+ return ordering;
+ }
+ }
+ left_len.cmp(&right_len)
+ }))
+ }
+ DataType::Struct(_) => {
+ let array = array.as_struct();
+ let fields = array
+ .columns()
+ .iter()
+ .map(|field| collated_comparator(field, collations,
unicode_version))
+ .collect::<Result<Vec<_>>>()?;
+ Ok(nulls_first(array.nulls().cloned(), move |left, right| {
+ fields
+ .iter()
+ .map(|compare| compare(left, right))
+ .find(|&ordering| ordering != Ordering::Equal)
+ .unwrap_or(Ordering::Equal)
+ }))
+ }
+ DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
+ if !collations.as_slice().is_empty() =>
+ {
+ let collation = collations.next().unwrap();
+ Ok(match array.data_type() {
+ DataType::Utf8 => string_comparator(
+ array.as_string::<i32>().clone(),
+ *collation,
+ unicode_version,
+ ),
+ DataType::LargeUtf8 => string_comparator(
+ array.as_string::<i64>().clone(),
+ *collation,
+ unicode_version,
+ ),
+ _ => string_comparator(array.as_string_view().clone(),
*collation, unicode_version),
+ })
+ }
+ _ => spark_comparator(array.as_ref(), array.as_ref()),
Review Comment:
The Scala side and this function each walk the element type depth-first, and
a string leaf picks its collation by position. Nothing checks that the two
walks agree. This arm does not consume an entry. So a string-bearing subtree it
does not recognize, such as a `LargeList`, a `FixedSizeList` or a
dictionary-encoded string (all of which `spark_comparator` accepts), would
compare bytewise and shift every later leaf onto the wrong collation without an
error. Comet does not produce those types here today, but a later change could
introduce that drift with no failing test. Could this arm return an error when
the type contains a string while collations remain?
##########
spark/src/test/resources/sql-tests/expressions/array/array_extrema_collation.sql:
##########
@@ -41,17 +48,39 @@ SELECT id, array_min(array(a, b)), array_max(array(a, b)),
FROM test_array_extrema_collation
-- Case-insensitive ordering differs from binary ordering for 'a' and 'B'.
-query expect_dispatch(array_min,array_max)
+-- Non-ASCII ties run in both input orders and must retain the first original
value.
+query expect_native(array_min,array_max)
SELECT id, array_min(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS
STRING COLLATE UTF8_LCASE))),
array_max(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS STRING
COLLATE UTF8_LCASE)))
FROM test_array_extrema_collation
--- Collation detection and Spark comparison also recurse through struct fields
and arrays.
-query expect_dispatch(array_min,array_max)
+-- Collation comparison also recurses through struct fields and arrays. The
second
+-- string field uses binary ordering, independently of the first field's
collation.
+query expect_native(array_min,array_max)
SELECT id, array_min(array(
- named_struct('s', array(CAST(a AS STRING COLLATE UTF8_LCASE)), 'f',
x),
- named_struct('s', array(CAST(b AS STRING COLLATE UTF8_LCASE)), 'f',
y))),
+ named_struct('s', array(CAST(a AS STRING COLLATE UTF8_LCASE)),
'binary', a, 'f', x),
Review Comment:
This query puts the collated field ahead of the binary one, so it would not
catch a `UTF8_BINARY` leaf that fails to consume its entry. Could you add a
variant with the binary field first, for example `named_struct('binary', a,
's', array(CAST(a AS STRING COLLATE UTF8_LCASE)), 'f', x)`? Then the positional
mapping is guarded from both directions.
##########
native/spark-expr/src/array_funcs/array_extrema.rs:
##########
@@ -156,9 +243,14 @@ where
result.finish()
}
-fn nested_extrema(array: &ListArray, is_min: bool) -> Result<ArrayRef> {
+fn nested_extrema(
+ array: &ListArray,
+ is_min: bool,
+ mut collations: std::slice::Iter<Utf8Collation>,
+ unicode_version: u32,
+) -> Result<ArrayRef> {
let values = array.values();
- let compare = spark_comparator(values.as_ref(), values.as_ref())?;
+ let compare = collated_comparator(values, &mut collations,
unicode_version)?;
Review Comment:
Related to the fallback arm below, could this check that `collations` is
empty once the comparator is built and return an error otherwise? Leftover
entries mean the native walk and `stringCollations` on the Scala side
disagreed, and today that would go unnoticed.
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -256,16 +256,57 @@ object CometArrayIntersect
private object ArrayExtremaSupport extends CometTypeShim {
val incompatReason: String =
- "Array extrema use binary string ordering for non-UTF8_BINARY collations "
+
+ "Array extrema support UTF8_BINARY, UTF8_LCASE, and UTF8_LCASE_RTRIM
natively; " +
+ "UTF8_LCASE requires Unicode 16 or 17. Other collations use binary
ordering when opted in " +
"(https://github.com/apache/datafusion-comet/issues/4496)."
- def getSupportLevel(elementType: DataType): SupportLevel = {
- if (hasNonDefaultStringCollation(elementType)) {
- // The dispatcher runs Spark's own comparison with the original
collation IDs, including
- // strings nested in arrays or structs. Keep the native bytewise
comparison opt-in only.
- Incompatible(Some(incompatReason))
- } else {
+ val binaryRtrimReason: String =
+ "UTF8_BINARY_RTRIM array extrema use Spark's comparator."
+
+ private def stringCollations(dt: DataType): Seq[String] = dt match {
+ case stringType: StringType => Seq(stringCollationName(stringType))
+ case ArrayType(elementType, _) => stringCollations(elementType)
+ case StructType(fields) => fields.toSeq.flatMap(f =>
stringCollations(f.dataType))
+ case _ => Seq.empty
+ }
+
+ private def supportsCollations(collations: Seq[String]): Boolean =
collations.forall {
+ case "UTF8_BINARY" | "UTF8_BINARY_RTRIM" => true
+ case "UTF8_LCASE" | "UTF8_LCASE_RTRIM" =>
+ collationUnicodeVersion == 16 || collationUnicodeVersion == 17
+ case _ => false
+ }
+
+ def getSupportLevel(elementType: DataType): SupportLevel =
+ // Keep this path on the dispatcher until its native performance
regression is resolved.
+ // Use Unsupported so allowIncompatible cannot silently change its
comparison semantics.
+ if (stringCollations(elementType).contains("UTF8_BINARY_RTRIM")) {
+ Unsupported(Some(binaryRtrimReason))
Review Comment:
The comment above says `Unsupported` keeps `allowIncompatible` from silently
changing the comparison semantics. With `Incompatible` and opt-in, though,
`convert` would send `UTF8_BINARY_RTRIM`, and the native comparator handles it
correctly, so the semantics would not change. The real reason is the
performance regression, and the comment should say that. This branch also makes
the native RTRIM support unreachable. That covers the `UTF8_BINARY_RTRIM` case
in `supportsCollations`, `Utf8Collation::BinaryRtrim` and its mapping in
`with_collations`, and the unit tests that exercise them. Could you either
remove those pieces, or keep them and open a tracking issue for the regression?
Then this branch, the comment and `binaryRtrimReason` could link to it.
`binaryRtrimReason` also ends up in the generated compatibility guide as an
unsupported reason, and "use Spark's comparator" alone does not tell a reader
why.
--
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]
