peter-toth commented on code in PR #57815: URL: https://github.com/apache/spark/pull/57815#discussion_r4483287992
########## sql/core/src/main/scala/org/apache/spark/sql/execution/window/DistinctWindowFunctionFrame.scala: ########## @@ -0,0 +1,578 @@ +/* + * 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 + // The window driver rejects partitions larger than Int.MaxValue rows before preparing frames. + // Use an exact conversion as a defensive check if that guard is bypassed. + processor.initialize(Math.toIntExact(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( Review Comment: **Finding 14.** For a frame that ends at `UNBOUNDED FOLLOWING`, the order in which unique values reach the processor depends on the dedup path: - **Hash path:** the destructive `BytesToBytesMap` iterator walks the data pages in append order, so values come in first-occurrence order. The growing frame keeps this order on both paths. - **Fallback:** after the key threshold, a failed `append`, or for a key that isn't binary-stable (e.g. `UTF8_LCASE`), `consumeDistinctRows` returns them in key order. Order-sensitive aggregates then disagree across the two paths. I ran this on `dc14878` with `o(id, x)` = `(1,'z'), (2,'a'), (3,'m'), (4,'z'), (5,'b'), (6,'q')`. Here `w` is `ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING`, and `wr` is the same with `RANGE`: | expression | default | `spark.sql.windowExec.distinct.hash.fallbackThreshold=1` | |---|---|---| | `collect_list(DISTINCT x) OVER w` | `[z, a, m, b, q]` | `[a, b, m, q, z]` | | `first(DISTINCT x) OVER wr` | `z` | `a` | | `collect_list(DISTINCT x COLLATE UTF8_LCASE) OVER w` | `[a, b, m, q, z]` | `[a, b, m, q, z]` | For comparison, `first(x) OVER wr` is `z`, and `array_distinct(collect_list(x) OVER w)` is `[z, a, m, b, q]` on both paths. The `first` case uses RANGE because `OptimizeWindowFunctions` rewrites `first` to `nth_value` on a ROWS frame. The fix replays the fallback through the event path the growing frame already uses, and leaves the hash path alone: ```scala override protected def processDistinctRows( firstVisibleRows: FirstVisibleRowsBuilder): Unit = { if (firstVisibleRows.isUsingMap) { // The map yields keys in insertion order, which is first-occurrence order. firstVisibleRows.foreachDistinctRow((key, _) => processor.update(key)) } else { // The fallback sorter yields keys in key order, so restore first-occurrence order. prepareOrderedEvents(firstVisibleRows) drainAllEvents() } } ``` It also needs `def isUsingMap: Boolean = usingMap` on `FirstVisibleRowsBuilder`, and this next to `updateProcessor`: ```scala protected final def drainAllEvents(): Unit = { while (nextEventIndex != Int.MaxValue) { processor.update(eventIterator.getValue) loadNextEvent() } } ``` With it, all three expressions return first-occurrence order on both paths. `SQLWindowFunctionSuite`, `DistinctWindowFunctionFrameSuite` and `DataFrameWindowFunctionsSuite` still pass. The `first(DISTINCT x)` query at `fallbackThreshold=1` and the collated `collect_list` at default both fail on `dc14878`, so either would pin the fix in a test. The fallback now reads `inputIndex`. So the finding 13 follow-up should drop only `firstVisibleIndex`, not the whole position row as I suggested there. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactoryBase.scala: ########## @@ -243,13 +260,69 @@ trait WindowEvaluatorFactoryBase { MutableProjection.create(expressions, schema), aggFilters) } + lazy val distinctAggregateExpressions = + functions.map(_.asInstanceOf[AggregateExpression]) + lazy val originalDistinctExpressions = + distinctChildren(distinctAggregateExpressions.head) + lazy val normalizedDistinctExpressions = + originalDistinctExpressions.map(NormalizeFloatingNumbers.normalize) + lazy val distinctInputAttributes = + normalizedDistinctExpressions.zipWithIndex.map { case (expression, index) => + AttributeReference( + s"windowDistinctValue$index", + expression.dataType, + expression.nullable)() + } + lazy val rewrittenDistinctFunctions: Array[Expression] = + distinctAggregateExpressions.map { ae => + val distinctColumnAttributeLookup = Utils.toMap( + distinctChildren(ae).map(_.canonicalized), + distinctInputAttributes) + ae.aggregateFunction.transformDown { + case expression => + distinctColumnAttributeLookup.getOrElse(expression.canonicalized, expression) + }.asInstanceOf[AggregateFunction] + } + lazy val distinctProcessor = AggregateProcessor( Review Comment: **Finding 15.** `distinctProcessor` is a `lazy val`, and so is `windowFrameExpressionFactoryPairs` on the factory. So every frame built from one factory object shares one `AggregateProcessor` and its buffer. Each task deserializes its own factory, so this only shows when one task evaluates the same window RDD twice. `zip` does that: ```scala val df = sql("""SELECT id, count(DISTINCT v) OVER (ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS d, count(v) OVER (ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS n FROM VALUES (1, 'a'), (2, 'b'), (3, 'a'), (4, 'c') AS t(id, v)""") val rdd = df.rdd rdd.zip(rdd).collect() ``` On `dc14878`, `df.collect()` gives `d` = `1, 2, 2, 3`, but the two sides of the zip give `1, 2, 2, 4` and `1, 3, 3, 5`. `n` is right on both sides, because the non-distinct branches call `def processor` once per frame. `rdd.zip(rdd.map(identity))` gives the same wrong values. With `def`, both sides return `1, 2, 2, 3`. The window suites still pass with it: ```suggestion def distinctProcessor = AggregateProcessor( ``` A test could compare both sides of `rdd.zip(rdd)` with `df.collect()`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/window/DistinctWindowFunctionFrame.scala: ########## @@ -0,0 +1,578 @@ +/* + * 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 + // The window driver rejects partitions larger than Int.MaxValue rows before preparing frames. + // Use an exact conversion as a defensive check if that guard is bypassed. + processor.initialize(Math.toIntExact(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) Review Comment: **Finding 16.** For a growing frame, `prepareOrderedEvents` calls `buildEvents()` for every window partition. On the hash path this builds a fresh `UnsafeKVExternalSorter` just to re-sort the map's entries. Each construction generates code (`SortPrefixUtils.createPrefixGenerator`, `GenerateOrdering.create`) and allocates the in-memory sorter's pointer array. On many small partitions that is most of the work. I measured it on `dc14878` with `spark.memory.debugFill=false` (sbt turns it on for tests), AQE off, over 400k rows with `PARTITION BY k ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`: | shape | `count(DISTINCT v)` | `size(collect_set(v))` | |---|---|---| | 200k partitions of 2 rows | 8.1 s | 0.12 s | | 20 partitions of 20k rows | 0.11 s | 0.12 s | The full-partition frame, which builds no event sorter, takes 0.32 s on the first shape. Stack samples during the slow run sit in `UnsafeKVExternalSorter.<init>`. On the hash path the sort is not needed. Candidates arrive in non-decreasing `(firstVisibleIndex, inputIndex)` order. The map keeps the first one per key, and `BytesToBytesMap`'s iterator walks the data pages in append order. So the map's entries are already in event order. I prototyped that. In the prototype, `prepareOrderedEvents` replays the map's entries while `usingMap` is still true, and builds the event sorter only after a fallback. That is the same split as the finding 14 fix. With it, the first shape takes 0.34 s. An assertion that the replayed positions only increase never fired. A differential check against `size(collect_set(...))` over 7 ROWS and RANGE frames, in hash, sorter and spill modes, still matches. The prototype held copies on the heap, so a real version would keep the map alive until the partition's output is done. "distinct event sorter spill size is reported" would then need to force a fallback first. -- 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]
