Merge branch 'main' into df54

Author: mbutrovich

docs/source/contributor-guide/adding_a_new_operator.md

@@ -553,8 +553,14 @@ For operators that run in the JVM:
 Example pattern from `CometExecRule.scala`:
 
 ```scala
-case s: ShuffleExchangeExec if nativeShuffleSupported(s) =>
-  CometShuffleExchangeExec(s, shuffleType = CometNativeShuffle)
+case s: ShuffleExchangeExec =>
+  CometShuffleExchangeExec.shuffleSupported(s) match {
+    case Some(CometNativeShuffle) =>
+      CometShuffleExchangeExec(s, shuffleType = CometNativeShuffle)
+    case Some(CometColumnarShuffle) =>
+      CometShuffleExchangeExec(s, shuffleType = CometColumnarShuffle)
+    case None => s
+  }
 ```
 
 ## Common Patterns and Helpers

docs/source/user-guide/latest/compatibility.md

@@ -58,6 +58,12 @@ Expressions that are not 100% Spark-compatible will fall back to Spark by defaul
 `spark.comet.expression.EXPRNAME.allowIncompatible=true`, where `EXPRNAME` is the Spark expression class name. See
 the [Comet Supported Expressions Guide](expressions.md) for more information on this configuration setting.
 
+### Aggregate Expressions
+
+- **CollectSet**: Comet deduplicates NaN values (treats `NaN == NaN`) while Spark treats each NaN as a distinct value.
+  When `spark.comet.exec.strictFloatingPoint=true`, `collect_set` on floating-point types falls back to Spark unless
+  `spark.comet.expression.CollectSet.allowIncompatible=true` is set.
+
 ### Array Expressions
 
 - **ArraysOverlap**: Inconsistent behavior when arrays contain NULL values.

docs/source/user-guide/latest/expressions.md

@@ -203,6 +203,7 @@ Expressions that are not Spark-compatible will fall back to Spark by default and
 | BitXorAgg     |            | Yes                       |                                                                  |
 | BoolAnd       | `bool_and` | Yes                       |                                                                  |
 | BoolOr        | `bool_or`  | Yes                       |                                                                  |
+| CollectSet    |            | No                        | NaN dedup differs from Spark. See compatibility guide.           |
 | Corr          |            | Yes                       |                                                                  |
 | Count         |            | Yes                       |                                                                  |
 | CovPopulation |            | Yes                       |                                                                  |

docs/spark_expressions_support.md

@@ -33,7 +33,7 @@
 - [x] bool_and
 - [x] bool_or
 - [ ] collect_list
-- [ ] collect_set
+- [x] collect_set
 - [ ] corr
 - [x] count
 - [x] count_if

native/core/src/execution/planner.rs

@@ -70,6 +70,7 @@ use datafusion_comet_spark_expr::{
     create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle,
     BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode, SumInteger, ToCsv,
 };
+use datafusion_spark::function::aggregate::collect::SparkCollectSet;
 use iceberg::expr::Bind;
 
 use crate::execution::operators::ExecutionError::GeneralError;
@@ -2259,6 +2260,11 @@ impl PhysicalPlanner {
                 ));
                 Self::create_aggr_func_expr("bloom_filter_agg", schema, vec![child], func)
             }
+            AggExprStruct::CollectSet(expr) => {
+                let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?;
+                let func = AggregateUDF::new_from_impl(SparkCollectSet::new());
+                Self::create_aggr_func_expr("collect_set", schema, vec![child], func)
+            }
         }
     }
 

native/proto/src/proto/expr.proto

@@ -140,6 +140,7 @@ message AggExpr {
     Stddev stddev = 14;
     Correlation correlation = 15;
     BloomFilterAgg bloomFilterAgg = 16;
+    CollectSet collectSet = 17;
   }
 
   // Optional filter expression for SQL FILTER (WHERE ...) clause.
@@ -248,6 +249,11 @@ message BloomFilterAgg {
   DataType datatype = 4;
 }
 
+message CollectSet {
+  Expr child = 1;
+  DataType datatype = 2;
+}
+
 enum EvalMode {
   LEGACY = 0;
   TRY = 1;

spark/src/main/scala/org/apache/comet/CometFallback.scala

@@ -200,22 +200,29 @@ object CometSparkSessionExtensions extends Logging {
   }
 
   /**
-   * Attaches explain information to a TreeNode, rolling up the corresponding information tags
-   * from any child nodes. For now, we are using this to attach the reasons why certain Spark
-   * operators or expressions are disabled.
+   * Record a fallback reason on a `TreeNode` (a Spark operator or expression) explaining why
+   * Comet cannot accelerate it. Reasons recorded here are surfaced in extended explain output
+   * (see `ExtendedExplainInfo`) and, when `COMET_LOG_FALLBACK_REASONS` is enabled, logged as
+   * warnings. The reasons are also rolled up from child nodes so that the operator that remains
+   * in the Spark plan carries the reasons from its converted-away subtree.
+   *
+   * Call this in any code path where Comet decides not to convert a given node - serde `convert`
+   * methods returning `None`, unsupported data types, disabled configs, etc. Do not use this for
+   * informational messages that are not fallback reasons: anything tagged here is treated by the
+   * rules as a signal that the node falls back to Spark.
    *
    * @param node
-   *   The node to attach the explain information to. Typically a SparkPlan
+   *   The Spark operator or expression that is falling back to Spark.
    * @param info
-   *   Information text. Optional, may be null or empty. If not provided, then only information
-   *   from child nodes will be included.
+   *   The fallback reason. Optional, may be null or empty - pass empty only when the call is used
+   *   purely to roll up reasons from `exprs`.
    * @param exprs
-   *   Child nodes. Information attached in these nodes will be be included in the information
-   *   attached to @node
+   *   Child nodes whose own fallback reasons should be rolled up into `node`. Pass the
+   *   sub-expressions or child operators whose failure caused `node` to fall back.
    * @tparam T
-   *   The type of the TreeNode. Typically SparkPlan, AggregateExpression, or Expression
+   *   The type of the TreeNode. Typically `SparkPlan`, `AggregateExpression`, or `Expression`.
    * @return
-   *   The node with information (if any) attached
+   *   `node` with fallback reasons attached (as a side effect on its tag map).
    */
   def withInfo[T <: TreeNode[_]](node: T, info: String, exprs: T*): T = {
     // support existing approach of passing in multiple infos in a newline-delimited string
@@ -228,22 +235,24 @@ object CometSparkSessionExtensions extends Logging {
   }
 
   /**
-   * Attaches explain information to a TreeNode, rolling up the corresponding information tags
-   * from any child nodes. For now, we are using this to attach the reasons why certain Spark
-   * operators or expressions are disabled.
+   * Record one or more fallback reasons on a `TreeNode` and roll up reasons from any child nodes.
+   * This is the set-valued form of [[withInfo]]; see that overload for the full contract.
+   *
+   * Reasons are accumulated (never overwritten) on the node's `EXTENSION_INFO` tag and are
+   * surfaced in extended explain output. When `COMET_LOG_FALLBACK_REASONS` is enabled, each new
+   * reason is also emitted as a warning.
    *
    * @param node
-   *   The node to attach the explain information to. Typically a SparkPlan
+   *   The Spark operator or expression that is falling back to Spark.
    * @param info
-   *   Information text. May contain zero or more strings. If not provided, then only information
-   *   from child nodes will be included.
+   *   The fallback reasons for this node. May be empty when the call is used purely to roll up
+   *   child reasons.
    * @param exprs
-   *   Child nodes. Information attached in these nodes will be be included in the information
-   *   attached to @node
+   *   Child nodes whose own fallback reasons should be rolled up into `node`.
    * @tparam T
-   *   The type of the TreeNode. Typically SparkPlan, AggregateExpression, or Expression
+   *   The type of the TreeNode. Typically `SparkPlan`, `AggregateExpression`, or `Expression`.
    * @return
-   *   The node with information (if any) attached
+   *   `node` with fallback reasons attached (as a side effect on its tag map).
    */
   def withInfos[T <: TreeNode[_]](node: T, info: Set[String], exprs: T*): T = {
     if (CometConf.COMET_LOG_FALLBACK_REASONS.get()) {
@@ -259,25 +268,27 @@ object CometSparkSessionExtensions extends Logging {
   }
 
   /**
-   * Attaches explain information to a TreeNode, rolling up the corresponding information tags
-   * from any child nodes
+   * Roll up fallback reasons from `exprs` onto `node` without adding a new reason of its own. Use
+   * this when a parent operator is itself falling back and wants to preserve the reasons recorded
+   * on its child expressions/operators so they appear together in explain output.
    *
    * @param node
-   *   The node to attach the explain information to. Typically a SparkPlan
+   *   The parent operator or expression falling back to Spark.
    * @param exprs
-   *   Child nodes. Information attached in these nodes will be be included in the information
-   *   attached to @node
+   *   Child nodes whose fallback reasons should be aggregated onto `node`.
    * @tparam T
-   *   The type of the TreeNode. Typically SparkPlan, AggregateExpression, or Expression
+   *   The type of the TreeNode. Typically `SparkPlan`, `AggregateExpression`, or `Expression`.
    * @return
-   *   The node with information (if any) attached
+   *   `node` with the rolled-up reasons attached (as a side effect on its tag map).
    */
   def withInfo[T <: TreeNode[_]](node: T, exprs: T*): T = {
     withInfos(node, Set.empty, exprs: _*)
   }
 
   /**
-   * Checks whether a TreeNode has any explain information attached
+   * True if any fallback reason has been recorded on `node` (via [[withInfo]] / [[withInfos]]).
+   * Callers that need to short-circuit when a prior rule pass has already decided a node falls
+   * back can use this as the sticky signal.
    */
   def hasExplainInfo(node: TreeNode[_]): Boolean = {
     node.getTagValue(CometExplainInfo.EXTENSION_INFO).exists(_.nonEmpty)

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -98,17 +98,18 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
   private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
 
   private def applyCometShuffle(plan: SparkPlan): SparkPlan = {
-    plan.transformUp {
-      case s: ShuffleExchangeExec if CometShuffleExchangeExec.nativeShuffleSupported(s) =>
-        // Switch to use Decimal128 regardless of precision, since Arrow native execution
-        // doesn't support Decimal32 and Decimal64 yet.
-        conf.setConfString(CometConf.COMET_USE_DECIMAL_128.key, "true")
-        CometShuffleExchangeExec(s, shuffleType = CometNativeShuffle)
-
-      case s: ShuffleExchangeExec if CometShuffleExchangeExec.columnarShuffleSupported(s) =>
-        // Columnar shuffle for regular Spark operators (not Comet) and Comet operators
-        // (if configured)
-        CometShuffleExchangeExec(s, shuffleType = CometColumnarShuffle)
+    plan.transformUp { case s: ShuffleExchangeExec =>
+      CometShuffleExchangeExec.shuffleSupported(s) match {
+        case Some(CometNativeShuffle) =>
+          // Switch to use Decimal128 regardless of precision, since Arrow native execution
+          // doesn't support Decimal32 and Decimal64 yet.
+          conf.setConfString(CometConf.COMET_USE_DECIMAL_128.key, "true")
+          CometShuffleExchangeExec(s, shuffleType = CometNativeShuffle)
+        case Some(CometColumnarShuffle) =>
+          CometShuffleExchangeExec(s, shuffleType = CometColumnarShuffle)
+        case None =>
+          s
+      }
     }
   }
 

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -20,7 +20,6 @@
 package org.apache.comet.serde
 
 import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Descending, NullsFirst, NullsLast, SortOrder}
-import org.apache.spark.sql.types._
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withInfo
@@ -30,19 +29,8 @@ object CometSortOrder extends CometExpressionSerde[SortOrder] {
 
   override def getSupportLevel(expr: SortOrder): SupportLevel = {
 
-    def containsFloatingPoint(dt: DataType): Boolean = {
-      dt match {
-        case DataTypes.FloatType | DataTypes.DoubleType => true
-        case ArrayType(elementType, _) => containsFloatingPoint(elementType)
-        case StructType(fields) => fields.exists(f => containsFloatingPoint(f.dataType))
-        case MapType(keyType, valueType, _) =>
-          containsFloatingPoint(keyType) || containsFloatingPoint(valueType)
-        case _ => false
-      }
-    }
-
     if (CometConf.COMET_EXEC_STRICT_FLOATING_POINT.get() &&
-      containsFloatingPoint(expr.child.dataType)) {
+      SupportLevel.containsFloatingPoint(expr.child.dataType)) {
       // https://github.com/apache/datafusion-comet/issues/2626
       Incompatible(
         Some(