feat: fall back to Spark for collated predicate operands

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -218,8 +218,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
       classOf[Chr] -> CometScalarFunction("char"),
       classOf[ConcatWs] -> CometConcatWs,
       classOf[Concat] -> CometConcat,
-      classOf[Contains] -> CometScalarFunction("contains"),
-      classOf[EndsWith] -> CometScalarFunction("ends_with"),
+      classOf[Contains] -> CometContains,
+      classOf[EndsWith] -> CometEndsWith,
       classOf[GetJsonObject] -> CometGetJsonObject,
       classOf[InitCap] -> CometInitCap,
       classOf[Length] -> CometLength,
@@ -233,7 +233,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
       classOf[RegExpReplace] -> CometRegExpReplace,
       classOf[Reverse] -> CometReverse,
       classOf[RLike] -> CometRLike,
-      classOf[StartsWith] -> CometScalarFunction("starts_with"),
+      classOf[StartsWith] -> CometStartsWith,
       classOf[StringInstr] -> CometScalarFunction("instr"),
       classOf[StringRepeat] -> CometStringRepeat,
       classOf[StringReplace] -> CometStringReplace,

spark/src/main/scala/org/apache/comet/serde/maps.scala

@@ -166,12 +166,17 @@ object CometMapFromEntries
   }
 }
 
-object CometStrToMap extends CometScalarFunction[StringToMap]("str_to_map") with CometTypeShim {
+object CometStrToMap
+    extends CometScalarFunction[StringToMap]("str_to_map")
+    with CometTypeShim
+    with CodegenDispatchFallback {
 
   // Spark 4.1.1+ honours spark.sql.legacy.truncateForEmptyRegexSplit by truncating trailing
   // empty entries from the split result. Comet's native str_to_map always behaves as if the flag
-  // were false, so it is incompatible when legacy truncation is enabled. Read by string key so it
-  // resolves on older Spark versions where the config is not registered.
+  // were false. When the flag is true, mark this Incompatible so the CodegenDispatchFallback
+  // trait routes the expression through the JVM codegen dispatcher (Spark's own doGenCode inside
+  // the Comet kernel) rather than falling the entire projection back to Spark. Read by string
+  // key so it resolves on older Spark versions where the config is not registered.
   private val legacyTruncateConfig = "spark.sql.legacy.truncateForEmptyRegexSplit"
 
   private val legacyTruncateReason =

spark/src/main/scala/org/apache/comet/serde/predicates.scala

@@ -21,7 +21,7 @@ package org.apache.comet.serde
 
 import scala.jdk.CollectionConverters._
 
-import org.apache.spark.sql.catalyst.expressions.{And, Attribute, EqualNullSafe, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, InSet, IsNaN, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not, Or}
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BinaryExpression, EqualNullSafe, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, InSet, IsNaN, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not, Or}
 import org.apache.spark.sql.types.BooleanType
 
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
@@ -34,26 +34,34 @@ object CometNot extends CometExpressionSerde[Not] {
       inputs: Seq[Attribute],
       binding: Boolean): Option[ExprOuterClass.Expr] = {
 
+    import ComparisonUtils.hasCollatedOperand
+
     expr.child match {
-      case expr: EqualTo =>
+      case inner: EqualTo if !hasCollatedOperand(inner.left, inner.right) =>
         createBinaryExpr(
-          expr,
-          expr.left,
-          expr.right,
+          inner,
+          inner.left,
+          inner.right,
           inputs,
           binding,
           (builder, binaryExpr) => builder.setNeq(binaryExpr))
-      case expr: EqualNullSafe =>
+      case inner: EqualNullSafe if !hasCollatedOperand(inner.left, inner.right) =>
         createBinaryExpr(
-          expr,
-          expr.left,
-          expr.right,
+          inner,
+          inner.left,
+          inner.right,
           inputs,
           binding,
           (builder, binaryExpr) => builder.setNeqNullSafe(binaryExpr))
-      case expr: In =>
-        ComparisonUtils.in(expr, expr.value, expr.list, inputs, binding, negate = true)
+      case inner: In if !hasCollatedOperand((inner.value +: inner.list): _*) =>
+        ComparisonUtils.in(inner, inner.value, inner.list, inputs, binding, negate = true)
       case _ =>
+        // Includes the collated variants of EqualTo / EqualNullSafe / In above: fall through so
+        // the child expression's own serde is consulted, which now returns `Unsupported` for
+        // non-UTF8_BINARY operands (see `CometEqualTo.getSupportLevel` and siblings). That makes
+        // `exprToProtoInternal` return None for the child, which cascades this Not to None and
+        // falls the enclosing operator back to Spark — the only way to honour collation-aware
+        // (in)equality without a native path.
         createUnaryExpr(
           expr,
           expr.child,
@@ -103,7 +111,7 @@ object CometOr extends CometExpressionSerde[Or] {
   }
 }
 
-object CometEqualTo extends CometExpressionSerde[EqualTo] {
+object CometEqualTo extends CollationAwareBinaryPredicate[EqualTo] {
   override def convert(
       expr: EqualTo,
       inputs: Seq[Attribute],
@@ -118,7 +126,7 @@ object CometEqualTo extends CometExpressionSerde[EqualTo] {
   }
 }
 
-object CometEqualNullSafe extends CometExpressionSerde[EqualNullSafe] {
+object CometEqualNullSafe extends CollationAwareBinaryPredicate[EqualNullSafe] {
   override def convert(
       expr: EqualNullSafe,
       inputs: Seq[Attribute],
@@ -133,7 +141,7 @@ object CometEqualNullSafe extends CometExpressionSerde[EqualNullSafe] {
   }
 }
 
-object CometGreaterThan extends CometExpressionSerde[GreaterThan] {
+object CometGreaterThan extends CollationAwareBinaryPredicate[GreaterThan] {
   override def convert(
       expr: GreaterThan,
       inputs: Seq[Attribute],
@@ -148,7 +156,7 @@ object CometGreaterThan extends CometExpressionSerde[GreaterThan] {
   }
 }
 
-object CometGreaterThanOrEqual extends CometExpressionSerde[GreaterThanOrEqual] {
+object CometGreaterThanOrEqual extends CollationAwareBinaryPredicate[GreaterThanOrEqual] {
   override def convert(
       expr: GreaterThanOrEqual,
       inputs: Seq[Attribute],
@@ -163,7 +171,7 @@ object CometGreaterThanOrEqual extends CometExpressionSerde[GreaterThanOrEqual]
   }
 }
 
-object CometLessThan extends CometExpressionSerde[LessThan] {
+object CometLessThan extends CollationAwareBinaryPredicate[LessThan] {
   override def convert(
       expr: LessThan,
       inputs: Seq[Attribute],
@@ -178,7 +186,7 @@ object CometLessThan extends CometExpressionSerde[LessThan] {
   }
 }
 
-object CometLessThanOrEqual extends CometExpressionSerde[LessThanOrEqual] {
+object CometLessThanOrEqual extends CollationAwareBinaryPredicate[LessThanOrEqual] {
   override def convert(
       expr: LessThanOrEqual,
       inputs: Seq[Attribute],
@@ -233,7 +241,14 @@ object CometIsNaN extends CometExpressionSerde[IsNaN] {
   }
 }
 
-object CometIn extends CometExpressionSerde[In] {
+object CometIn extends CometExpressionSerde[In] with CodegenDispatchFallback {
+
+  override def getSupportLevel(expr: In): SupportLevel =
+    ComparisonUtils.collationSupportLevel("In", (expr.value +: expr.list): _*)
+
+  override def getUnsupportedReasons(): Seq[String] =
+    Seq(ComparisonUtils.nonDefaultCollationDocReason)
+
   override def convert(
       expr: In,
       inputs: Seq[Attribute],
@@ -242,7 +257,14 @@ object CometIn extends CometExpressionSerde[In] {
   }
 }
 
-object CometInSet extends CometExpressionSerde[InSet] {
+object CometInSet extends CometExpressionSerde[InSet] with CodegenDispatchFallback {
+
+  override def getSupportLevel(expr: InSet): SupportLevel =
+    ComparisonUtils.collationSupportLevel("InSet", expr.child)
+
+  override def getUnsupportedReasons(): Seq[String] =
+    Seq(ComparisonUtils.nonDefaultCollationDocReason)
+
   override def convert(
       expr: InSet,
       inputs: Seq[Attribute],
@@ -257,8 +279,52 @@ object CometInSet extends CometExpressionSerde[InSet] {
   }
 }
 
+/**
+ * Mixin for serdes of binary predicates whose native kernel compares raw bytes: any operand
+ * carrying a non-UTF8_BINARY collation is unsupported natively. Mixing in
+ * `CodegenDispatchFallback` routes these through the JVM codegen dispatcher (Spark's own
+ * `doGenCode`), so the collation-aware evaluation runs inline in the Comet pipeline instead of
+ * falling the whole operator back to Spark.
+ */
+trait CollationAwareBinaryPredicate[T <: BinaryExpression]
+    extends CometExpressionSerde[T]
+    with CodegenDispatchFallback {
+  override def getSupportLevel(expr: T): SupportLevel =
+    ComparisonUtils.collationSupportLevel(expr.prettyName, expr.left, expr.right)
+
+  override def getUnsupportedReasons(): Seq[String] =
+    Seq(ComparisonUtils.nonDefaultCollationDocReason)
+}
+
 object ComparisonUtils {
 
+  // Comet's native equality/ordering/hashing compare raw bytes, so any predicate operand carrying
+  // a non-UTF8_BINARY collation (Spark 4+) would produce wrong answers on the native path — e.g.
+  // `'a' = 'A'` under `UNICODE_CI` returns true in Spark but false byte-wise. Every binary
+  // comparison and `In`/`InSet` serde routes its `getSupportLevel` through here, marking the case
+  // `Unsupported`; `CodegenDispatchFallback` then routes those cases through Spark's own
+  // `doGenCode` inside the Comet pipeline. `hasNonDefaultStringCollation` walks nested types too,
+  // so collated strings inside array/map/struct operands are also caught.
+  def nonDefaultCollationReason(exprName: String): String =
+    s"$exprName does not support non-UTF8_BINARY collated operands; " +
+      "native comparison is byte-wise and cannot honour collation semantics."
+
+  // Doc-friendly variant of `nonDefaultCollationReason` for `getUnsupportedReasons`. The compat
+  // guide already scopes the reason to a specific expression, so no per-expr name is needed.
+  val nonDefaultCollationDocReason: String =
+    "Non-UTF8_BINARY collated operands are routed through the JVM codegen dispatcher " +
+      "(Spark's own `doGenCode`) because native comparison is byte-wise."
+
+  def hasCollatedOperand(operands: Expression*): Boolean =
+    operands.exists(op => hasNonDefaultStringCollation(op.dataType))
+
+  def collationSupportLevel(exprName: String, operands: Expression*): SupportLevel =
+    if (hasCollatedOperand(operands: _*)) {
+      Unsupported(Some(nonDefaultCollationReason(exprName)))
+    } else {
+      Compatible()
+    }
+
   def in(
       expr: Expression,
       value: Expression,

spark/src/main/scala/org/apache/comet/serde/strings.scala

@@ -19,7 +19,7 @@
 
 package org.apache.comet.serde
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64, BitLength, Cast, Concat, ConcatWs, Elt, Empty2Null, Expression, FindInSet, FormatNumber, FormatString, GetJsonObject, InitCap, Left, Length, Levenshtein, Like, Literal, Lower, Mask, OctetLength, Overlay, RegExpExtract, RegExpExtractAll, RegExpInStr, RegExpReplace, Right, RLike, SoundEx, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringSplit, StringTranslate, Substring, SubstringIndex, ToCharacter, ToNumber, TryToNumber, UnBase64, Upper}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64, BitLength, Cast, Concat, ConcatWs, Contains, Elt, Empty2Null, EndsWith, Expression, FindInSet, FormatNumber, FormatString, GetJsonObject, InitCap, Left, Length, Levenshtein, Like, Literal, Lower, Mask, OctetLength, Overlay, RegExpExtract, RegExpExtractAll, RegExpInStr, RegExpReplace, Right, RLike, SoundEx, StartsWith, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringSplit, StringTranslate, Substring, SubstringIndex, ToCharacter, ToNumber, TryToNumber, UnBase64, Upper}
 import org.apache.spark.sql.types.{BinaryType, DataTypes, LongType, StringType}
 
 import org.apache.comet.CometConf
@@ -288,13 +288,21 @@ object CometConcatWs extends CometExpressionSerde[ConcatWs] {
   }
 }
 
-object CometLike extends CometExpressionSerde[Like] {
+object CometLike extends CometExpressionSerde[Like] with CodegenDispatchFallback {
+
+  private val customEscapeReason =
+    "LIKE with a custom escape character (only `\\` is supported natively)"
+
+  override def getUnsupportedReasons(): Seq[String] =
+    Seq(customEscapeReason, ComparisonUtils.nonDefaultCollationDocReason)
 
   override def getSupportLevel(expr: Like): SupportLevel = {
-    if (expr.escapeChar == '\\') {
-      Compatible()
-    } else {
+    if (ComparisonUtils.hasCollatedOperand(expr.left, expr.right)) {
+      Unsupported(Some(ComparisonUtils.nonDefaultCollationReason("Like")))
+    } else if (expr.escapeChar != '\\') {
       Unsupported(Some(s"custom escape character ${expr.escapeChar} not supported in LIKE"))
+    } else {
+      Compatible()
     }
   }
 
@@ -309,6 +317,24 @@ object CometLike extends CometExpressionSerde[Like] {
   }
 }
 
+/**
+ * Serdes for `Contains` / `StartsWith` / `EndsWith` that reject non-UTF8_BINARY collated operands
+ * and otherwise delegate to the generic `contains` / `starts_with` / `ends_with` scalar-function
+ * bridge. The native kernels compare raw bytes and cannot honour case- or accent-insensitive
+ * collations, so a collated operand must fall back to Spark.
+ */
+object CometContains
+    extends CometScalarFunction[Contains]("contains")
+    with CollationAwareBinaryPredicate[Contains]
+
+object CometStartsWith
+    extends CometScalarFunction[StartsWith]("starts_with")
+    with CollationAwareBinaryPredicate[StartsWith]
+
+object CometEndsWith
+    extends CometScalarFunction[EndsWith]("ends_with")
+    with CollationAwareBinaryPredicate[EndsWith]
+
 /**
  * `rlike` runs Spark's own implementation through the codegen dispatcher by default, for
  * byte-exact results. The native (rust) regexp engine is faster but has different semantics from

spark/src/test/resources/sql-tests/expressions/map/str_to_map_collation.sql

@@ -17,17 +17,16 @@
 
 -- MinSparkVersion: 4.0
 
--- Spark 4.0+ applies collation-aware delimiter matching in str_to_map. Comet's native
--- implementation matches raw delimiter bytes, so non-UTF8_BINARY collations must fall back.
+-- Spark 4.0+ applies collation-aware delimiter matching in str_to_map.
 
 -- collated input string
-query expect_fallback(`str_to_map` does not support non-UTF8_BINARY collations)
+query
 SELECT str_to_map('a:1,b:2' COLLATE UTF8_LCASE, ',', ':')
 
 -- collated pair delimiter
-query expect_fallback(`str_to_map` does not support non-UTF8_BINARY collations)
+query
 SELECT str_to_map('aX1,bX2', 'x' COLLATE UTF8_LCASE, ':')
 
 -- collated key-value delimiter
-query expect_fallback(`str_to_map` does not support non-UTF8_BINARY collations)
+query
 SELECT str_to_map('aX1,bX2', ',', 'x' COLLATE UTF8_LCASE)

spark/src/test/resources/sql-tests/expressions/map/str_to_map_legacy_truncate.sql

@@ -15,19 +15,22 @@
 -- specific language governing permissions and limitations
 -- under the License.
 
--- Tests that str_to_map falls back to Spark when
--- spark.sql.legacy.truncateForEmptyRegexSplit is enabled. In legacy mode Spark truncates trailing
--- empty entries from the split result, which Comet's native str_to_map does not honour.
+-- Tests that str_to_map routes through the JVM codegen dispatcher when
+-- spark.sql.legacy.truncateForEmptyRegexSplit is enabled. In legacy mode Spark truncates
+-- trailing empty entries from the split result, which Comet's native str_to_map does not
+-- honour. `CometStrToMap` marks the expression Incompatible when the flag is on and mixes in
+-- [[CodegenDispatchFallback]], so the projection stays native (Spark's own `doGenCode` runs
+-- inside the Comet kernel) while producing Spark-exact results.
 -- See https://github.com/apache/datafusion-comet/issues/4477
 
 -- Config: spark.sql.legacy.truncateForEmptyRegexSplit=true
 
--- trailing pair delimiter: legacy mode truncates the trailing empty entry, so Comet must fall
--- back to Spark
-query expect_fallback(truncateForEmptyRegexSplit)
+-- trailing pair delimiter: legacy mode truncates the trailing empty entry; Comet delegates to
+-- the codegen dispatcher.
+query
 SELECT str_to_map('a:1,b:2,', ',', ':')
 
--- column input also falls back
+-- column input is also handled via the codegen dispatcher
 statement
 CREATE TABLE test_str_to_map_legacy(s STRING, pair_delim STRING, key_value_delim STRING) USING parquet
 
@@ -37,5 +40,5 @@ INSERT INTO test_str_to_map_legacy VALUES
   ('x:1;y:2;', ';', ':'),
   (NULL, ',', ':')
 
-query expect_fallback(truncateForEmptyRegexSplit)
+query
 SELECT str_to_map(s, pair_delim, key_value_delim) FROM test_str_to_map_legacy