MaxGekk commented on code in PR #56808:
URL: https://github.com/apache/spark/pull/56808#discussion_r3482875357
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala:
##########
@@ -241,10 +242,48 @@ class DataFramePivotSuite extends SharedSparkSession {
)
}
- test("pivot with null should not throw NPE") {
+ // SPARK-19882: null pivot keys must not NPE; also checks empty count
buckets return 0.
+ test("pivot count with null pivot value returns 0 for empty buckets") {
checkAnswer(
Seq(Tuple1(None),
Tuple1(Some(1))).toDF("a").groupBy($"a").pivot("a").count(),
- Row(null, 1, null) :: Row(1, null, 1) :: Nil)
+ Row(null, 1, 0) :: Row(1, 0, 1) :: Nil)
+ }
+
+ // Empty-bucket result coverage is in the pivot.sql golden file. The tests
below cover what it
+ // cannot: output-schema nullability, the ANSI runtime error path, and the
empty-bucket flag.
+
+ test("pivot count produces non-nullable column schema") {
+ val df = Seq((1, "a")).toDF("id", "cat")
+ .groupBy("id").pivot("cat", Seq("a", "b")).count()
+ assert(!df.schema("a").nullable, s"expected 'a' column to be non-nullable:
${df.schema}")
+ assert(!df.schema("b").nullable, s"expected 'b' column to be non-nullable:
${df.schema}")
+ }
+
+ test("pivot integral division of counts throws on an empty bucket under
ANSI") {
Review Comment:
The composite-aggregate coverage here is all never-NULL-on-non-empty shapes
(`count(v1) div count(v2)` throws under ANSI; the bare-count default), so the
suite passes even with the composite `Coalesce` issue above present. Worth
adding a case that exercises a composite which is NULL on a *non-empty* bucket
— e.g. `nullif(count(v), N)` over a bucket with exactly `N` matching rows
asserting `NULL`, and `if(count(v1) > 0, sum(v2), 0)` over a non-empty
all-NULL-`v2` bucket asserting `NULL` — comparing against the slow path /
flag-off.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PivotTransformer.scala:
##########
@@ -185,6 +201,42 @@ object PivotTransformer extends AliasHelper with
SQLConfHelper {
}
}
+ /**
+ * Empty-input default for a pivot aggregate to coalesce into its extracted
value, or `None` to
+ * leave the value unchanged. The fast path's [[PivotFirst]] leaves an
unmatched pivot category's
+ * slot unset, so the caller wraps the result in a [[Coalesce]] to recover
the value the slow path
+ * produces on an empty bucket (`count` -> 0; `sum`/`avg`/`min`/`max` ->
NULL).
+ *
+ * Returned unevaluated so later constant folding defers a default that
throws under ANSI (e.g.
+ * `count(v1) / count(v2)` -> `0 / 0`) to runtime, where [[Coalesce]] only
evaluates it for
+ * actually-empty buckets -- matching the slow path. Mirrors
+ * `RewriteCorrelatedScalarSubquery.evalAggExprOnZeroTups`.
+ */
+ private def aggregateEmptyInputDefault(aggregate: Expression):
Option[Expression] = {
+ trimAliases(aggregate) match {
+ // Bare aggregate: use its published default result (count -> 0,
sum/avg/min/max -> None).
+ case AggregateExpression(aggregateFunction, _, _, _, _) =>
+ aggregateFunction.defaultResult
+ // Composite over aggregate(s): substitute each aggregate/attribute with
its empty-input
+ // value. Return None for a non-foldable default or a literal NULL
(nothing to coalesce in). A
+ // default that only folds to NULL (e.g. sum(x) + 1) is still returned;
its Coalesce evaluates
+ // to NULL on an empty bucket, which is the correct result.
+ case other =>
+ val default = other.transform {
+ case AggregateExpression(aggregateFunction, _, _, _, _) =>
+ aggregateFunction.defaultResult.getOrElse(
+ Literal.create(null, aggregateFunction.dataType))
+ case attribute: AttributeReference =>
+ Literal.create(null, attribute.dataType)
+ }
+ default match {
+ case _ if !trimAliases(default).foldable => None
+ case Literal(null, _) => None
Review Comment:
Nit: this guard only catches a *syntactic* `Literal(null)`, but a default
that *evaluates* to null without being a literal — e.g. `CEIL(null)` (from
`CEIL(sum(x))`) or `(null + cast(1 as bigint))` (from `sum(x) + 1`) — slips
through and yields `Some(default)`, so a `Coalesce(extractedValue,
<expr→null>)` wrapper is emitted (you can see it in the analyzer-results
golden, e.g. `coalesce(__pivot_CEIL(sum…), CEIL(null))`). It's harmless to
results (`Coalesce(x, null) ≡ x`, and nullability is unchanged) but adds dead
plan nodes and doesn't match the doc's "Return None for … a literal NULL".
Folding the default first would catch these: `default.foldable &&
default.eval() == null => None`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PivotTransformer.scala:
##########
@@ -185,6 +201,42 @@ object PivotTransformer extends AliasHelper with
SQLConfHelper {
}
}
+ /**
+ * Empty-input default for a pivot aggregate to coalesce into its extracted
value, or `None` to
+ * leave the value unchanged. The fast path's [[PivotFirst]] leaves an
unmatched pivot category's
+ * slot unset, so the caller wraps the result in a [[Coalesce]] to recover
the value the slow path
+ * produces on an empty bucket (`count` -> 0; `sum`/`avg`/`min`/`max` ->
NULL).
+ *
+ * Returned unevaluated so later constant folding defers a default that
throws under ANSI (e.g.
+ * `count(v1) / count(v2)` -> `0 / 0`) to runtime, where [[Coalesce]] only
evaluates it for
+ * actually-empty buckets -- matching the slow path. Mirrors
+ * `RewriteCorrelatedScalarSubquery.evalAggExprOnZeroTups`.
+ */
+ private def aggregateEmptyInputDefault(aggregate: Expression):
Option[Expression] = {
+ trimAliases(aggregate) match {
+ // Bare aggregate: use its published default result (count -> 0,
sum/avg/min/max -> None).
+ case AggregateExpression(aggregateFunction, _, _, _, _) =>
+ aggregateFunction.defaultResult
+ // Composite over aggregate(s): substitute each aggregate/attribute with
its empty-input
+ // value. Return None for a non-foldable default or a literal NULL
(nothing to coalesce in). A
+ // default that only folds to NULL (e.g. sum(x) + 1) is still returned;
its Coalesce evaluates
+ // to NULL on an empty bucket, which is the correct result.
+ case other =>
+ val default = other.transform {
Review Comment:
The composite branch can return a non-null `default` for an aggregate that
is NULL-capable on a *non-empty* bucket, and the caller then wraps it in
`Coalesce(extractedValue, default)`. The problem is that `PivotFirst.update`
stores a value only `if (value != null)` and keeps no presence bit, so a buffer
slot is NULL iff the pivot category is **absent** *or*
**present-but-the-value-is-NULL** — `ExtractValue` can't distinguish them, and
the `Coalesce` defaults both. That's correct for bare count-family aggregates
(never NULL on a non-empty bucket), but wrong for composites over a
NULL-capable sub-aggregate:
- `PIVOT(nullif(count(v), 5) ...)` — a bucket with exactly 5 matching
non-null rows: slow path = `nullif(5,5)` = NULL, fast path = `Coalesce(NULL,
nullif(0,5)=0)` = `0`.
- `PIVOT(if(count(v1) > 0, sum(v2), 0) ...)` — a non-empty bucket with
all-NULL `v2`: slow path = NULL, fast path = `0`.
Since the flag defaults to `true`, these are live wrong results, and the
pivoted column is also declared non-nullable while it can be NULL. The simplest
safe fix is to apply the empty-default only for bare aggregates whose
`defaultResult` is non-null (count-family, where a NULL slot provably means
empty) and leave NULL-capable composites to NULL / the slow path; a fuller fix
would add a per-slot presence signal to `PivotFirst`.
--
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]
