[PR] [SPARK-60124][SQL] Fix RewriteDistinctAggregates failure on ListAgg with DISTINCT and FILTER [spark]

["zahed1994 (via GitHub)" <gi...@apache.org>]

zahed1994 opened a new pull request, #59460:
URL: https://github.com/apache/spark/pull/59460

   ### What changes were proposed in this pull request?
   In `RewriteDistinctAggregates`, aggregate functions with `SortOrder` in their children (specifically `ListAgg` with `WITHIN GROUP (ORDER BY ...)`) caused a `ClassCastException` during distinct rewrite when combined with `FILTER` or when regular `ListAgg` is evaluated alongside other distinct aggregates:
   
   ```
   java.lang.ClassCastException: class org.apache.spark.sql.catalyst.expressions.AttributeReference cannot be cast to class org.apache.spark.sql.catalyst.expressions.SortOrder
       at org.apache.spark.sql.catalyst.expressions.aggregate.ListAgg.withNewChildrenInternal(collect.scala:1051)
       at org.apache.spark.sql.catalyst.optimizer.RewriteDistinctAggregates.patchAggregateFunctionChildren(RewriteDistinctAggregates.scala:262)
   ```
   
   This occurred because:
   1. `SortOrder` extends `Expression with Unevaluable`, but Catalyst treats `orderExpressions` as part of `ListAgg.children`.
   2. `RewriteDistinctAggregates` treated `SortOrder` directly as an unfoldable child expression, causing `Expand` to project `SortOrder` (an `Unevaluable` expression).
   3. `patchAggregateFunctionChildren` mapped `af.children` with `attrs(c).getOrElse(c)`, replacing the `SortOrder` itself with an `AttributeReference`. When `ListAgg.withNewChildrenInternal` cast `newChildren.drop(2)` back to `SortOrder`, it threw `ClassCastException`.
   
   This PR fixes the issue by:
   1. Introducing an `evalChildren` helper that unpacks `SortOrder.children` to operate on evaluable expressions when determining foldability, grouping sets, and projecting attributes.
   2. In `patchAggregateFunctionChildren`, unwrapping `SortOrder`, replacing only its child expressions with rewritten attributes, and preserving the `SortOrder` wrapper.
   
   ### Why are the changes needed?
   Queries that use `LISTAGG(DISTINCT ...) WITHIN GROUP (ORDER BY ...) FILTER (...)` or regular `LISTAGG(...) WITHIN GROUP (ORDER BY ...)` alongside distinct aggregates fail during query planning with an internal `ClassCastException`.
   
   ### Does this PR introduce _any_ user-facing change?
   No.
   
   ### How was this patch tested?
   - Added unit tests in `RewriteDistinctAggregatesSuite`.
   - Added end-to-end execution tests in `DataFrameAggregateSuite`.
   
   ### Was this patch authored or co-authored using generative AI tooling?
   Yes, Generated-by: Google Antigravity.
   


-- 
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-60124][SQL] Fix RewriteDistinctAggregates failure on ListAgg with DISTINCT and FILTER [spark]

["HyukjinKwon (via GitHub)" <gi...@apache.org>]

HyukjinKwon commented on code in PR #59460:
URL: https://github.com/apache/spark/pull/59460#discussion_r4281025876


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala:
##########
@@ -230,7 +237,8 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
 
     // Extract distinct aggregate expressions.
     val distinctAggGroups = aggExpressions.filter(_.isDistinct).groupBy { e =>
-        val unfoldableChildren = ExpressionSet(e.aggregateFunction.children.filter(!_.foldable))
+        val unfoldableChildren = ExpressionSet(
+          evalChildren(e.aggregateFunction.children).filter(!_.foldable))

Review Comment:
   **Non-blocking (P2):** With `evalChildren` in the group key, `count(DISTINCT s)` and `listagg(DISTINCT s, '|') WITHIN GROUP (ORDER BY s)` both key to `{s}`. The same happens for two `listagg(DISTINCT s)` calls ordered ASC and DESC, or one with and one without `WITHIN GROUP`. Without a `FILTER`, `distinctAggGroups.size == 1` and `mustRewrite` is false, so the rule returns the `Aggregate` unchanged.
   
   `SparkStrategies.Aggregation` still computes `distinctAggChildSets` from the raw children (`ExpressionSet(ae.aggregateFunction.children.filterNot(_.foldable))`, SparkStrategies.scala:734-742). It sees `{s}` vs `{s, SortOrder(s)}` and throws `You hit a query analyzer bug`:
   
   ```sql
   SELECT count(DISTINCT s), listagg(DISTINCT s, '|') WITHIN GROUP (ORDER BY s) FROM t
   ```
   
   This query also failed before the PR, with the `ClassCastException`, so it is not a regression. But the failure only moves from the rewrite to the planner. The planner already unwraps `SortOrder` when it builds `distinctExpressions` (SparkStrategies.scala:768-773); only the sanity check disagrees with the rule's new grouping. With a `FILTER`, the same queries are rewritten and succeed.
   
   **Recommended change:** Compute the planner's distinctAggChildSets with SortOrder children replaced by their child expressions, the same view RewriteDistinctAggregates now uses for its group key and the strategy already uses for distinctExpressions. Add end-to-end aggregate tests for count(DISTINCT s) with listagg(DISTINCT s) WITHIN GROUP (ORDER BY s) and for ASC versus DESC listagg(DISTINCT s), both without FILTER.
   
   **Why this works:** Once both components agree on the distinct child set, every Aggregate the rule declines to rewrite has exactly one set as the planner computes it. The single-distinct path then plans it, deduplicating on s and evaluating each distinct aggregate, including every differently ordered listagg, over the deduplicated rows.
   
   **Scope:** Align the physical planner's distinct-aggregate sanity check with the rule's SortOrder-unwrapped distinct grouping, and cover the merged-group shape end to end.
   
   **Compatibility:** Aggregates whose distinct child sets differ after unwrapping must still be rejected by the sanity check and handled by RewriteDistinctAggregates.
   
   **Risks:** The sanity check could become more permissive than the single-distinct path can execute. The traced canonicalized transformDown substitution handles SortOrder children, and end-to-end tests must confirm the merged and ASC/DESC shapes produce correct results.
   
   **Constraints:** Multiple genuinely different distinct child sets must still trigger the internal error, because RewriteDistinctAggregates rewrites them.
   
   **Success:** SELECT count(DISTINCT s), listagg(DISTINCT s, '|') WITHIN GROUP (ORDER BY s) without FILTER plans and returns the distinct count and the ordered concatenation of distinct values. Two listagg(DISTINCT s) calls ordered ASC and DESC in one query without FILTER return the two oppositely ordered results. Aggregates whose distinct child sets differ after SortOrder unwrapping still fail the planner sanity check if they ever reach it.
   
   **Verification:**
   - **Regression:** End-to-end query combining count(DISTINCT s) with listagg(DISTINCT s) WITHIN GROUP (ORDER BY s), without FILTER, returns correct values instead of an internal error.
   - **Behavior:** End-to-end query with listagg(DISTINCT s) ordered ASC and DESC in the same SELECT returns both orderings correctly.
   - **Compatibility:** Existing distinct-aggregate planning and RewriteDistinctAggregates coverage keeps passing, including single-distinct listagg golden queries.



##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala:
##########
@@ -5461,7 +5461,33 @@ class DataFrameAggregateSuite extends SharedSparkSession
     assert(outType.precision == 24 && outType.scale == 6,
       s"expected DecimalType(24, 6) from un-rewritten AVG, got $outType")
   }
+
+  test("SPARK-60124: listagg(DISTINCT ...) with FILTER") {
+    val df = Seq(
+      ("b", 1), ("a", 1), ("a", 1), (null.asInstanceOf[String], 1), ("c", 0)
+    ).toDF("s", "keep")
+
+    checkAnswer(
+      df.selectExpr("listagg(DISTINCT s, '|') WITHIN GROUP (ORDER BY s) FILTER (WHERE keep = 1)"),
+      Seq(Row("a|b"))
+    )
+    checkAnswer(
+      df.selectExpr(
+        "count(DISTINCT keep)",
+        "listagg(DISTINCT s, '|') WITHIN GROUP (ORDER BY s) FILTER (WHERE keep = 1)"
+      ),
+      Seq(Row(2, "a|b"))
+    )
+    checkAnswer(
+      df.selectExpr(
+        "count(DISTINCT keep)",
+        "listagg(s, '|') WITHIN GROUP (ORDER BY s)"

Review Comment:
   **Nit (P3):** This case never reaches `RewriteDistinctAggregates`. With a single distinct aggregate, no `FILTER`, and a non-foldable `keep`, `mayNeedtoRewrite` returns false, and the planner's single-distinct path plans the query exactly as it did before this PR. The assertion passes with or without the fix. As a result, the regular-aggregate half of the change (line 340 and the `SortOrder` branch of `patchAggregateFunctionChildren`) has no end-to-end result check; only the catalyst test's `assert(rewrite.resolved)` reaches it. Adding a second distinct group, e.g. `count(DISTINCT s)` with expected `Row(2, 3, "a|a|b|c")`, would force the rewrite.
   
   See **Shared repair plan 1** in the review body.



##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala:
##########
@@ -5461,7 +5461,33 @@ class DataFrameAggregateSuite extends SharedSparkSession
     assert(outType.precision == 24 && outType.scale == 6,
       s"expected DecimalType(24, 6) from un-rewritten AVG, got $outType")
   }
+
+  test("SPARK-60124: listagg(DISTINCT ...) with FILTER") {
+    val df = Seq(
+      ("b", 1), ("a", 1), ("a", 1), (null.asInstanceOf[String], 1), ("c", 0)
+    ).toDF("s", "keep")
+
+    checkAnswer(
+      df.selectExpr("listagg(DISTINCT s, '|') WITHIN GROUP (ORDER BY s) FILTER (WHERE keep = 1)"),

Review Comment:
   **Nit (P3):** Every new test orders by the aggregated column itself (`ORDER BY s` over `s`), so the order child and the aggregated child collapse to the same `Expand` column and attribute. A bug in how the new `SortOrder` branch remaps the order child, for example mapping it to the child's attribute, would not change any assertion.
   
   The distinguishing shape is reachable by default. With `spark.sql.listagg.allowDistinctCastWithOrder.enabled` (default true), `listagg(DISTINCT i, ',') WITHIN GROUP (ORDER BY i) FILTER (...)` over an int column becomes child `CAST(i AS STRING)` with `SortOrder(i)` and takes this rewrite with key `{CAST(i AS STRING), i}`. Values such as 9 and 10 would expose string versus numeric ordering. A regular `listagg(s) WITHIN GROUP (ORDER BY o)` next to two distinct groups is the same gap on the regular path.
   
   See **Shared repair plan 1** in the review body.



-- 
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