namanjain24-sudo commented on PR #25289: URL: https://github.com/apache/datafusion/pull/25289#issuecomment-5681874559
@jayzhan211 thanks, you're right. My earlier timings were on copies of the loops outside DataFusion, so this time I measured inside it: a criterion bench that calls the three UDFs, comparing `main` (133111f), this PR before the change, and this PR with your suggestions. I ran it twice. Rows with values in [-1, 1] never reach the rescue path. They run at 0.92x to 1.03x of `main` for `array_distance` and `cosine_distance`, and `array_normalize` is still 0.55x to 0.65x at dim 128 and 1536. The regression was in rows that reach the rescue path without benefiting from it (time over `main`, before → after your changes): | input | dim 4 | dim 128 | dim 1536 | | --- | --- | --- | --- | | `cosine_distance`, one vector all zero | 1.35x → 1.11x | 3.45x → 1.13x | 4.67x → 1.17x | | `array_distance`, identical vectors | 0.94x → 0.94–0.98x | 1.51x → 1.00–1.04x | 2.35x → 1.19–1.20x | | `array_normalize`, zero vector | 1.18x → 1.20–1.21x | 2.39x → 1.63–1.65x | 1.98x → 0.99–1.29x | A range means the two runs differed. For `array_normalize` at dim 1536, `main` itself moved between runs. I've pushed both suggestions as b3fcbe1: - `norm_scale` checks finiteness once on the max, so the loop vectorizes. - `cosine_distance` scans a vector only when its own sum is out of range, and recomputes only if there is something to scale. The commit also includes your doc fix for scaled subnormals and the bit-pattern assertions. To check that no result moved, I ran the previous and new kernels on about 3.7 million vector pairs. They covered lengths from 1 to 1536, magnitudes from subnormal up to 1e300, and zero, NaN, infinite and one-side-only inputs. All three functions are bit-identical to the previous commit, and `cosine_distance` also matches a version that always rescales both vectors. The length-based threshold from @aoto-tech's review is unchanged, so rows around 1e-100 still never reach the rescue path (0 scans in about 87,000 random rows). On the inputs they flagged in the issue, NaN, infinite, zero and identical vectors give the same bits as `main`. A dot product that overflows and cancels (`[1e200, 1e200]` with `[1e200, -1e200]`) gives `1` like the always-rescaled version, where `main` returns `NaN`. Two small additions to your comment: - Rows whose sum is not zero but below `len * 2^-1012` also take the rescue path. Those are the rows that actually underflowed, so they are the ones it is for. - A vector left unscaled can also be one containing an infinity, where `norm_scale` returns `None`, so its sum of squares is not always finite. The result is still `NaN` exactly as on `main`, so I worded the code comment around the sum being in range. The bench isn't part of this PR. I can add it as `benches/vector_functions.rs` if that would help. -- 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]
