peter-toth commented on code in PR #57815: URL: https://github.com/apache/spark/pull/57815#discussion_r3881680243
########## sql/core/src/main/scala/org/apache/spark/sql/execution/window/DistinctWindowFunctionFrame.scala: ########## @@ -0,0 +1,576 @@ +/* + * 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.memory.SparkOutOfMemoryError +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. If the map cannot be allocated, or if 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. + * If map construction fails, inputs go directly to the sorter. 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() + try { + new BytesToBytesMap(taskMemoryManager, 64, taskMemoryManager.pageSizeBytes()) + } catch { + case _: SparkOutOfMemoryError => null + } + } 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) + bufferUpdated = true + loadNextEvent() + } + if (bufferUpdated) { + processor.evaluate(target) + } + } + + override final def currentLowerBound(): Int = 0 + + private def closeEventResources(): Unit = { + if (eventSorter != null) { + spillSize.add(eventSorter.getSpillSize) + } + if (eventIterator != null) { + eventIterator.close() + eventIterator = null + } + if (eventSorter != null) { + eventSorter.cleanupResources() + eventSorter = null + } + } + + override final def close(): Unit = closeEventResources() +} + +/** + * Computes DISTINCT aggregates over the entire window partition. Since every output row has the + * same set of unique inputs, their consumption order is undefined and is not restored after hash + * deduplication or sorting by the DISTINCT key. + */ +private[window] final class UnboundedDistinctWindowFunctionFrame( + target: InternalRow, + processor: AggregateProcessor, + distinctExpressions: Seq[Expression], + filter: Option[Expression], + inputSchema: Seq[Attribute], + spillSize: SQLMetric, + hashFallbackThreshold: Int, + spillThreshold: Int, + spillSizeThreshold: Long) + extends DistinctWindowFunctionFrame( + target, + processor, + distinctExpressions, + filter, + inputSchema, + spillSize, + hashFallbackThreshold, + spillThreshold, + spillSizeThreshold) { + + private var partitionSize = 0 + + override protected def populateFirstVisibleRows( + rows: ExternalAppendOnlyUnsafeRowArray, + firstVisibleRows: FirstVisibleRowsBuilder): Unit = { + val iterator = rows.generateIterator() + var inputIndex = 0 + while (iterator.hasNext) { + addCandidate(iterator.next(), 0, inputIndex, firstVisibleRows) + inputIndex += 1 + } + } + + override protected def processDistinctRows( + firstVisibleRows: FirstVisibleRowsBuilder): Unit = { + firstVisibleRows.foreachDistinctRow((key, _) => processor.update(key)) Review Comment: **Finding 13.** This path discards the position, and `populateFirstVisibleRows` right above always passes `firstVisibleIndex = 0`, so the `(firstVisibleIndex, inputIndex)` value declared at `:62-64` is written for every candidate row of a full-partition frame and read by nothing. It costs on both dedup paths. Hash path: 24 bytes of value next to a 16-byte `BIGINT` key. I measured a `BytesToBytesMap` with 8-byte keys at `spark.buffer.pageSize=1m`, comparing the 24-byte position against a zero-length value - `getTotalMemoryConsumption` is 10,240 KiB vs 8,192 KiB at 100k keys (1.25x) and 38,912 KiB vs 29,696 KiB at 400k keys (1.31x). Now that `spark.sql.windowExec.distinct.hash.fallbackThreshold` defaults to `Int.MaxValue`, the map degrades when `Location.append` fails, i.e. on real memory exhaustion, so that 25-31% is exactly the margin that decides when `count(DISTINCT x) OVER ()` stops using the hash path. Sorter path: the value is inserted per *candidate row*, not per distinct key, and `consumeDistinctRows` (`:299-303`) runs an `isEarlier` comparison plus two `copy()` per record to pick a representative whose position nobody reads. Shape of the fix: make the position schema a `protected def` that this subclass overrides to `StructType(Nil)`, and have `consumeDistinctRows` keep the first row of each key group when it is empty - `isEarlier` reads `getInt(0)`, which is not valid on a zero-field row. The zero-length record path is already proven by the `count(DISTINCT 1)` empty-key case. Follow-up-only, nothing here blocks the merge. Reasonable to fold into the JIRA you are opening for finding 11. -- 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]
