peter-toth commented on code in PR #57815:
URL: https://github.com/apache/spark/pull/57815#discussion_r3758365266
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala:
##########
@@ -196,11 +197,380 @@ class SQLWindowFunctionSuite extends SharedSparkSession {
val e = intercept[AnalysisException] {
sql(
"""
- |select month, area, product, sum(distinct product + 1) over
(partition by 1 order by 2)
+ |select month, area, product, sum(distinct product + 1) over (
+ | partition by 1 order by 2 rows between current row and current
row)
|from windowData
""".stripMargin)
}
- assert(e.getMessage.contains("Distinct window functions are not
supported"))
+ assert(e.getMessage.contains("Unsupported DISTINCT window function"))
+ }
+ }
+
+ test("window function: distinct rejects unorderable inputs") {
+ val e = intercept[AnalysisException] {
+ sql("SELECT count(DISTINCT map('key', id)) OVER () FROM range(1)")
+ }
+ assert(e.getCondition === "DISTINCT_WINDOW_FUNCTION_UNSUPPORTED")
+ }
+
+ test("window function: distinct aggregates with an unbounded preceding
frame") {
+ val data = Seq(
+ (1, 0, 10, "a", 10),
+ (1, 1, 20, "a", 10),
+ (1, 2, 20, "b", 20),
+ (1, 3, 20, null.asInstanceOf[String], 30),
+ (1, 4, 30, "c", 30),
+ (2, 5, 5, "b", 5),
+ (2, 6, 5, "b", 5),
+ (2, 7, 6, "a", 6)
+ ).toDF("k", "id", "v", "x", "amount")
+
+ withTempView("distinctWindowData") {
+ data.createOrReplaceTempView("distinctWindowData")
+
+ checkAnswer(
+ sql(
+ """
+ |SELECT k, id,
+ | count(DISTINCT x) OVER (PARTITION BY k ORDER BY v) AS
range_count,
+ | count(DISTINCT x) OVER (
+ | PARTITION BY k ORDER BY v, id
+ | ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS
rows_count,
+ | count(DISTINCT x) OVER (
+ | PARTITION BY k ORDER BY v, id
+ | ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING) AS
preceding_count,
+ | count(DISTINCT x) OVER (
+ | PARTITION BY k ORDER BY v, id
+ | ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING) AS
following_count,
+ | count(DISTINCT x) OVER (PARTITION BY k) AS partition_count,
+ | sum(DISTINCT amount) OVER (PARTITION BY k ORDER BY v) AS
range_sum,
+ | avg(DISTINCT amount) OVER (PARTITION BY k ORDER BY v) AS
range_avg,
+ | sort_array(collect_list(DISTINCT amount) OVER (
+ | PARTITION BY k ORDER BY v)) AS range_values
+ |FROM distinctWindowData
+ """.stripMargin),
+ Seq(
+ Row(1, 0, 1L, 1L, 0L, 1L, 3L, 10L, 10.0, Seq(10)),
+ Row(1, 1, 2L, 1L, 1L, 2L, 3L, 60L, 20.0, Seq(10, 20, 30)),
+ Row(1, 2, 2L, 2L, 1L, 2L, 3L, 60L, 20.0, Seq(10, 20, 30)),
+ Row(1, 3, 2L, 2L, 2L, 3L, 3L, 60L, 20.0, Seq(10, 20, 30)),
+ Row(1, 4, 3L, 3L, 2L, 3L, 3L, 60L, 20.0, Seq(10, 20, 30)),
+ Row(2, 5, 1L, 1L, 0L, 1L, 2L, 5L, 5.0, Seq(5)),
+ Row(2, 6, 1L, 1L, 1L, 2L, 2L, 5L, 5.0, Seq(5)),
+ Row(2, 7, 2L, 2L, 1L, 2L, 2L, 11L, 5.5, Seq(5, 6))
+ ))
+ }
+ }
+
+ test("window function: count distinct with a range offset, filter, and
multiple columns") {
Review Comment:
**Finding 9.** The RANGE branch of `populateFirstVisibleRows`
(`sql/core/src/main/scala/org/apache/spark/sql/execution/window/DistinctWindowFunctionFrame.scala:501-517`)
is the only code in this PR that holds *two* live iterators over `rows` at the
same time, and `currentUpperBound()` can add a third on top of the evaluator's
own `bufferIterator`. It does work — `generateIterator` builds independent
readers and each `SpillableArrayIterator` owns its own `UnsafeRow` — but no
test runs it against a spilled buffer, so if that ever stops holding nothing
here will say so.
The specific gap is (spilled buffer x non-empty distinct key):
- this test and `"window function: distinct aggregates with an unbounded
preceding frame"` (`:216`) are the only RANGE cases, and neither sets
`WINDOW_EXEC_BUFFER_IN_MEMORY_THRESHOLD`, so the buffer stays on the
`ArrayBuffer` path and `generateIterator` returns `InMemoryBufferIterator`;
- `"window function: distinct works when the window input buffer spills"`
(`:558`) is the only test that does spill the buffer, and it uses
`count(DISTINCT 1)` — a zero-field key, so `distinctProjection` never reads a
byte out of the reused row;
- `"window function: distinct handles binary-unstable collation across
spills"` (`:410`) sets only `WINDOW_EXEC_BUFFER_SPILL_THRESHOLD`, which does
nothing until the in-memory threshold is crossed first, so that buffer doesn't
spill either.
Cheapest fix is one conf on this test, which already covers RANGE plus an
offset, a filter and a two-column key:
```scala
test("window function: count distinct with a range offset, filter, and
multiple columns") {
withSQLConf(WINDOW_EXEC_BUFFER_IN_MEMORY_THRESHOLD.key -> "1") {
// ... existing body
}
}
```
Same expected rows, but now read out of an `UnsafeExternalSorter`-backed
buffer through three concurrent iterators.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactoryBase.scala:
##########
@@ -169,6 +183,9 @@ trait WindowEvaluatorFactoryBase {
case e@WindowExpression(function, spec) =>
val frame =
spec.frameSpecification.asInstanceOf[SpecifiedWindowFrame]
function match {
+ case ae @ AggregateExpression(_, _, true, _, _)
+ if frame.lower == UnboundedPreceding =>
+ collect("DISTINCT_AGGREGATE", frame, e, ae)
Review Comment:
**Finding 8.** Before this PR the invariant here was total:
`WindowResolution.checkWindowFunction` rejected *every* distinct aggregate, so
the `case AggregateExpression(f, _, _, _, _)` on the next line could never see
one — and the scaladoc on `eligibleForSegTree` said so. Now it's conditional. A
distinct aggregate avoids that case only because the analyzer's frame predicate
and this one happen to agree.
They are written out separately, in two files, with nothing tying them
together:
- analyzer,
`sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala:171-180`:
`SpecifiedWindowFrame(_, UnboundedPreceding, _)` **and**
`RowOrdering.isOrderable` over the non-foldable children with `SortOrder`
unwrapped;
- here, `:186-187`: `frame.lower == UnboundedPreceding`.
If they ever drift the failure is silent and wrong, not loud:
`collect("AGGREGATE", frame, e, f)` passes `ae.aggregateFunction`, so
`isDistinct` is dropped on the floor and the query returns the non-distinct
answer. No exception, no plan difference, nothing a test would trip over.
The child extraction is duplicated the same way — `distinctChildren` at
`:150-155` and the inline copy at `WindowResolution.scala:174-177` — and the
two already differ, since only this one picked up `distinctBy(_.canonicalized)`
in the last commit. It happens not to matter (`isOrderable` gives the same
answer either way), but it is the same list defined twice.
Two small changes would close it. Make the impossible case say so:
```scala
case ae @ AggregateExpression(_, _, true, _, _) if frame.lower ==
UnboundedPreceding =>
collect("DISTINCT_AGGREGATE", frame, e, ae)
case AggregateExpression(_, _, true, _, _) =>
// Rejected by WindowResolution.checkWindowFunction; reaching here means
the two
// gates disagree, and silently dropping DISTINCT would give a wrong
answer.
throw SparkException.internalError(
s"DISTINCT is not supported for window frame ${frame.sql}")
case AggregateExpression(f, _, _, _, _) => collect("AGGREGATE", frame, e, f)
```
and hoist the child extraction into one helper both sides call, so the gate
and the dedup key are provably the same list.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/window/DistinctWindowFunctionFrame.scala:
##########
@@ -0,0 +1,548 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.window
+
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.UnsafeRowUtils
+import org.apache.spark.sql.execution.{ExternalAppendOnlyUnsafeRowArray,
UnsafeKVExternalSorter}
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.{DataType, IntegerType, StructField,
StructType}
+import org.apache.spark.unsafe.map.BytesToBytesMap
+
+/**
+ * Computes one or more equivalent DISTINCT aggregate expressions for a frame
whose lower bound is
+ * UNBOUNDED PRECEDING.
+ *
+ * Such a frame never removes rows: it either covers the entire partition or
only grows as output
+ * advances. For each qualifying input row this frame finds the first output
row whose upper bound
+ * contains it. A BytesToBytesMap removes duplicates while it remains below
configurable key-count
+ * and memory-size soft limits. When another new key arrives after either
limit has been reached, or
+ * an append fails, the frame permanently falls back to an external sorter for
the rest of the
+ * window partition. For a growing frame, that sorter finds the earliest event
for every key, then
+ * a second external sorter orders those events by (firstVisibleIndex,
inputIndex). The frame feeds
+ * each unique, normalized DISTINCT input to the aggregate processor when its
event becomes visible.
+ * A frame covering the entire partition consumes unique inputs directly
because their order is
+ * undefined.
+ */
+private[window] abstract class DistinctWindowFunctionFrame(
+ target: InternalRow,
+ processor: AggregateProcessor,
+ distinctExpressions: Seq[Expression],
+ filter: Option[Expression],
+ inputSchema: Seq[Attribute],
+ spillSize: SQLMetric,
+ hashFallbackThreshold: Int,
+ spillThreshold: Int,
+ spillSizeThreshold: Long)
+ extends WindowFunctionFrame with AutoCloseable {
+
+ private val distinctFields = distinctExpressions.zipWithIndex.map { case
(expression, index) =>
+ StructField(s"key$index", expression.dataType, expression.nullable)
+ }
+ private val distinctKeySchema = StructType(distinctFields)
+ private val distinctTypes = distinctKeySchema.map(_.dataType)
+ private val positionSchema = StructType(Seq(
+ StructField("firstVisibleIndex", IntegerType, nullable = false),
+ StructField("inputIndex", IntegerType, nullable = false)))
+
+ private val distinctProjection =
UnsafeProjection.create(distinctExpressions, inputSchema)
+ private val distinctOrdering =
RowOrdering.createNaturalAscendingOrdering(distinctTypes)
+ private val canUseHashDedup =
distinctTypes.forall(UnsafeRowUtils.isBinaryStable)
+ private val boundFilter = filter.map(Predicate.create(_, inputSchema))
+ private val positionInput =
+ new SpecificInternalRow(Seq(IntegerType, IntegerType))
+ private val positionProjection =
+ UnsafeProjection.create(Array[DataType](IntegerType, IntegerType))
+
+ private var eventSorter: UnsafeKVExternalSorter = _
+ private var eventIterator: UnsafeKVExternalSorter#KVSorterIterator = _
+ private var nextEventIndex = Int.MaxValue
+
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ =>
close()))
+
+ override final def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
+ closeEventResources()
+ nextEventIndex = Int.MaxValue
+ processor.initialize(rows.length)
+ val partitionIndex =
Option(TaskContext.get()).map(_.partitionId()).getOrElse(0)
+ boundFilter.foreach(_.initialize(partitionIndex))
+
+ val firstVisibleRows = new FirstVisibleRowsBuilder
+ try {
+ populateFirstVisibleRows(rows, firstVisibleRows)
+ processDistinctRows(firstVisibleRows)
+ spillSize.add(firstVisibleRows.getSpillSize)
+ } finally {
+ firstVisibleRows.close()
+ }
+ prepareFrame(rows)
+ }
+
+ protected def populateFirstVisibleRows(
+ rows: ExternalAppendOnlyUnsafeRowArray,
+ firstVisibleRows: FirstVisibleRowsBuilder): Unit
+
+ protected def processDistinctRows(firstVisibleRows:
FirstVisibleRowsBuilder): Unit
+
+ protected def prepareFrame(rows: ExternalAppendOnlyUnsafeRowArray): Unit
+
+ protected final def prepareOrderedEvents(firstVisibleRows:
FirstVisibleRowsBuilder): Unit = {
+ var newEventSorter: UnsafeKVExternalSorter = null
+ try {
+ newEventSorter = firstVisibleRows.buildEvents()
+ eventSorter = newEventSorter
+ newEventSorter = null
+ eventIterator = eventSorter.sortedIterator()
+ loadNextEvent()
+ } finally {
+ if (newEventSorter != null) {
+ newEventSorter.cleanupResources()
+ }
+ }
+ }
+
+ protected final def addCandidate(
+ row: InternalRow,
+ firstVisibleIndex: Int,
+ inputIndex: Int,
+ firstVisibleRows: FirstVisibleRowsBuilder): Unit = {
+ if (boundFilter.forall(_.eval(row))) {
+ positionInput.setInt(0, firstVisibleIndex)
+ positionInput.setInt(1, inputIndex)
+ val distinctRow = distinctProjection(row)
+ firstVisibleRows.add(distinctRow, positionProjection(positionInput))
+ }
+ }
+
+ /**
+ * Uses a bounded BytesToBytesMap to remove duplicates before they reach the
first external
+ * sorter. The map contains only the earliest row for each distinct key
because candidates arrive
+ * in non-decreasing (firstVisibleIndex, inputIndex) order.
+ *
+ * BytesToBytesMap compares raw UnsafeRow bytes, so binary-unstable keys use
the sorter directly.
+ * Once the map reaches either its key-count or memory-size soft limit and
another new key
+ * arrives, or it cannot append a record, all map entries are transferred to
one external sorter.
+ * The rest of this window partition goes directly to that sorter and never
returns to hash-based
+ * deduplication.
+ */
+ protected final class FirstVisibleRowsBuilder extends AutoCloseable {
+ private val map = if (canUseHashDedup) {
+ val taskMemoryManager = TaskContext.get().taskMemoryManager()
+ new BytesToBytesMap(taskMemoryManager, 64,
taskMemoryManager.pageSizeBytes())
+ } else {
+ null
+ }
+ private var usingMap = map != null
+ private var sorter: UnsafeKVExternalSorter = _
+
+ def add(key: UnsafeRow, value: UnsafeRow): Unit = {
+ if (!usingMap) {
+ insertIntoSorter(key, value)
+ return
+ }
+
+ val location = map.lookup(
+ key.getBaseObject,
+ key.getBaseOffset,
+ key.getSizeInBytes)
+ if (location.isDefined) {
+ return
+ }
+
+ if (map.numKeys() >= hashFallbackThreshold ||
+ map.getTotalMemoryConsumption >= spillSizeThreshold) {
+ switchToSorter()
+ insertIntoSorter(key, value)
+ } else if (!location.append(
+ key.getBaseObject,
+ key.getBaseOffset,
+ key.getSizeInBytes,
+ value.getBaseObject,
+ value.getBaseOffset,
+ value.getSizeInBytes)) {
+ switchToSorter()
+ insertIntoSorter(key, value)
+ }
+ }
+
+ def buildEvents(): UnsafeKVExternalSorter = {
+ var events: UnsafeKVExternalSorter = null
+ try {
+ if (usingMap) {
+ // Make the map spillable before allocating the event sorter. The
destructive iterator
+ // releases the hash array immediately and lets memory pressure
spill remaining pages.
+ val iterator = destructiveMapIterator()
+ events = newSorter(positionSchema, distinctKeySchema)
+ drainMap(iterator, (key, position) => emitEvent(events, position,
key))
+ } else {
+ events = newSorter(positionSchema, distinctKeySchema)
+ if (sorter != null) {
+ consumeDistinctRows(
+ sorter, (key, position) => emitEvent(events, position, key))
+ }
+ }
+ val result = events
+ events = null
+ result
+ } finally {
+ if (events != null) {
+ events.cleanupResources()
+ }
+ }
+ }
+
+ def foreachDistinctRow(consume: (UnsafeRow, UnsafeRow) => Unit): Unit = {
+ if (usingMap) {
+ drainMap(destructiveMapIterator(), consume)
+ } else if (sorter != null) {
+ consumeDistinctRows(sorter, consume)
+ }
+ }
+
+ /**
+ * Returns bytes spilled by the first, distinct-key sorter. The frame
accounts for the second,
+ * event-order sorter separately when closing its event resources. Spill
files created while a
+ * destructive BytesToBytesMap iterator releases the map's data pages are
not included in the
+ * window spill metric.
+ */
+ def getSpillSize: Long = if (sorter == null) 0L else sorter.getSpillSize
+
+ private def insertIntoSorter(key: UnsafeRow, value: UnsafeRow): Unit = {
+ if (sorter == null) {
+ sorter = newSorter(distinctKeySchema, positionSchema)
+ }
+ sorter.insertKV(key, value)
+ }
+
+ private def switchToSorter(): Unit = {
+ assert(sorter == null)
+ val iterator = destructiveMapIterator()
+ sorter = newSorter(distinctKeySchema, positionSchema)
+ drainMap(iterator, (key, position) => sorter.insertKV(key, position))
+ }
+
+ private def destructiveMapIterator(): BytesToBytesMap#MapIterator = {
+ assert(usingMap)
+ val iterator = map.destructiveIterator()
+ usingMap = false
+ iterator
+ }
+
+ private def drainMap(
+ iterator: BytesToBytesMap#MapIterator,
+ consume: (UnsafeRow, UnsafeRow) => Unit): Unit = {
+ val key = new UnsafeRow(distinctKeySchema.length)
+ val position = new UnsafeRow(positionSchema.length)
+ while (iterator.hasNext) {
+ val location = iterator.next()
+ key.pointTo(
+ location.getKeyBase,
+ location.getKeyOffset,
+ location.getKeyLength)
+ position.pointTo(
+ location.getValueBase,
+ location.getValueOffset,
+ location.getValueLength)
+ consume(key, position)
+ }
+ }
+
+ override def close(): Unit = {
+ if (sorter != null) {
+ sorter.cleanupResources()
+ sorter = null
+ }
+ if (map != null) {
+ map.free()
+ usingMap = false
+ }
+ }
+ }
+
+ private def consumeDistinctRows(
+ firstSorter: UnsafeKVExternalSorter,
+ consume: (UnsafeRow, UnsafeRow) => Unit): Unit = {
+ val iterator = firstSorter.sortedIterator()
+ var groupKey: UnsafeRow = null
+ var selectedKey: UnsafeRow = null
+ var selectedPosition: UnsafeRow = null
+ try {
+ while (iterator.next()) {
+ val key = iterator.getKey
+ val position = iterator.getValue
+ if (groupKey == null) {
+ groupKey = key.copy()
+ selectedKey = groupKey
+ selectedPosition = position.copy()
+ } else if (distinctOrdering.compare(groupKey, key) == 0) {
+ if (isEarlier(position, selectedPosition)) {
+ selectedKey = key.copy()
+ selectedPosition = position.copy()
+ }
+ } else {
+ consume(selectedKey, selectedPosition)
+ groupKey = key.copy()
+ selectedKey = groupKey
+ selectedPosition = position.copy()
+ }
+ }
+ if (groupKey != null) {
+ consume(selectedKey, selectedPosition)
+ }
+ } finally {
+ iterator.close()
+ }
+ }
+
+ private def isEarlier(left: InternalRow, right: InternalRow): Boolean = {
+ val leftVisibleIndex = left.getInt(0)
+ val rightVisibleIndex = right.getInt(0)
+ leftVisibleIndex < rightVisibleIndex ||
+ leftVisibleIndex == rightVisibleIndex && left.getInt(1) < right.getInt(1)
+ }
+
+ private def emitEvent(
+ events: UnsafeKVExternalSorter,
+ position: UnsafeRow,
+ distinctValue: UnsafeRow): Unit = {
+ events.insertKV(position, distinctValue)
+ }
+
+ private def newSorter(
+ keySchema: StructType,
+ valueSchema: StructType): UnsafeKVExternalSorter = {
+ val taskContext = TaskContext.get()
+ // The frame owns all of its sorters and registers one task completion
listener at construction.
+ UnsafeKVExternalSorter.createWithCallerOwnedLifecycle(
+ keySchema,
+ valueSchema,
+ SparkEnv.get.blockManager,
+ SparkEnv.get.serializerManager,
+ taskContext.taskMemoryManager().pageSizeBytes,
+ spillThreshold,
+ spillSizeThreshold)
+ }
+
+ private def loadNextEvent(): Unit = {
+ if (eventIterator != null && eventIterator.next()) {
+ nextEventIndex = eventIterator.getKey.getInt(0)
+ } else {
+ nextEventIndex = Int.MaxValue
+ }
+ }
+
+ protected final def updateProcessor(index: Int): Unit = {
+ var bufferUpdated = index == 0
+ while (nextEventIndex <= index) {
+ processor.update(eventIterator.getValue)
Review Comment:
You're right and I was wrong — withdrawing this finding.
The claim it rested on is false: `CodeGenerator.setColumn` does *not* fall
through to a bare `row.update` for these types, it has an explicit copy case
(`sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:1750-1752`),
with a comment saying exactly why, and `InternalRow.getWriter` does the same
for the interpreted path
(`sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/InternalRow.scala:209-214`).
So `updateProjection` copies into the processor's own `buffer` before any page
is freed, and `evaluate` reads that buffer, not the input row — the chain I
described breaks at the very first hop, at both hand-off sites. `BinaryType` is
in neither copy list but `UnsafeRow.getBinary` allocates a fresh array, and the
imperative aggregates go through `InternalRow.copyValue`, so those are covered
too. No `.copy()` needed here.
The `max(DISTINCT string)` / `first(DISTINCT string)` cases you added are
worth keeping anyway — they pin the behaviour so nobody has to re-derive this.
--
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]
