urosstan-db commented on code in PR #59615:
URL: https://github.com/apache/spark/pull/59615#discussion_r4474652717
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/vectorExpressions.scala:
##########
@@ -674,11 +676,74 @@ object VectorAvg {
newCount: Long): Unit = {
// Update running average: new_avg = old_avg + (new_value - old_avg) /
new_count
val invCount = 1.0f / newCount
+ if (dim < 4) {
+ updateRemainingElements(accBytes, inputArray, 0, dim, invCount)
+ } else {
+ updateElementsInBlocks(accBytes, inputArray, dim, invCount)
+ }
+ }
+
+ private def updateElementsInBlocks(
+ accBytes: Array[Byte],
+ inputArray: ArrayData,
+ dim: Int,
+ invCount: Float): Unit = {
var i = 0
+ while (i <= dim - 4) {
+ val off = Platform.BYTE_ARRAY_OFFSET + i.toLong * 4
+ val oldAvg0 = Platform.getFloat(accBytes, off)
+ val delta0 = inputArray.getFloat(i) - oldAvg0
+ val oldAvg1 = Platform.getFloat(accBytes, off + 4)
+ val delta1 = inputArray.getFloat(i + 1) - oldAvg1
+ val oldAvg2 = Platform.getFloat(accBytes, off + 8)
+ val delta2 = inputArray.getFloat(i + 2) - oldAvg2
+ val oldAvg3 = Platform.getFloat(accBytes, off + 12)
+ val delta3 = inputArray.getFloat(i + 3) - oldAvg3
+ // A finite sum guarantees finite deltas; sum overflow only triggers a
scalar fallback.
+ val deltaSum = (delta0 + delta1) + (delta2 + delta3)
+ if (!JFloat.isFinite(deltaSum)) {
+ updateRemainingElements(accBytes, inputArray, i, dim, invCount)
+ return
+ }
+ Platform.putFloat(accBytes, off, oldAvg0 + delta0 * invCount)
+ Platform.putFloat(accBytes, off + 4, oldAvg1 + delta1 * invCount)
+ Platform.putFloat(accBytes, off + 8, oldAvg2 + delta2 * invCount)
+ Platform.putFloat(accBytes, off + 12, oldAvg3 + delta3 * invCount)
+ i += 4
+ }
+ if (i < dim) {
+ updateRemainingElements(accBytes, inputArray, i, dim, invCount)
+ }
+ }
+
+ private def updateRemainingElements(
+ accBytes: Array[Byte],
+ inputArray: ArrayData,
+ start: Int,
+ dim: Int,
+ invCount: Float): Unit = {
+ var i = start
while (i < dim) {
val off = Platform.BYTE_ARRAY_OFFSET + i.toLong * 4
val oldAvg = Platform.getFloat(accBytes, off)
- Platform.putFloat(accBytes, off, oldAvg + (inputArray.getFloat(i) -
oldAvg) * invCount)
+ val value = inputArray.getFloat(i)
+ val delta = value - oldAvg
+ val newAvg = if (JFloat.isFinite(delta)) {
+ oldAvg + delta * invCount
+ } else if (JFloat.isFinite(oldAvg) && JFloat.isFinite(value)) {
+ // Scale before subtracting when opposite finite values overflow the
delta.
+ // For oldAvg = Float.MaxValue, value = -Float.MaxValue and invCount
= 0.5f,
+ // the delta is -Infinity, but scaling first gives the correct
average of 0.0f.
+ oldAvg + (value * invCount - oldAvg * invCount)
Review Comment:
Thanks for pointing this out. We reproduced it and tried a guard that falls
back when finite partial averages produce a non-finite weighted sum. In local
`VectorAvg.merge` microbenchmarks on Java 21, that candidate increased ordinary
finite-merge time by **106.2% at dimension 128** (0.568 -> 1.171 ns/element)
and **95.3% at dimension 1536** (0.376 -> 0.735 ns/element). These are kernel
measurements, not end-to-end query timings.
We dropped that candidate and are leaving the merge fix outside this PR.
This is a valid, pre-existing numerical edge case involving averages near
`Float.MaxValue` and partition counts whose independently rounded weights sum
to more than one. Given the common-path cost of the implementation we tried, we
prefer to handle it separately and look for a cheaper solution. The
measurements apply to that candidate; they do not establish that every possible
fix has the same cost.
The merge implementation remains unchanged, and the migration note is now
scoped specifically to updating running averages.
--
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]
