mbutrovich opened a new issue, #6781: URL: https://github.com/apache/datafusion-comet/issues/6781
### What is the problem the feature request solves? Two per-row functions divide twice where one division is enough. - [`pmod`](https://github.com/apache/datafusion-comet/blob/f9dd86d482d136426990186c2f1d42eef6a1ca59/native/shuffle/src/comet_partitioning.rs#L123-L129) maps each row's hash to a shuffle partition. When `hash % n` is negative, it computes `(r + n) % n`. In an aarch64 release build (rustc 1.98.1), both `%` operations compile to `sdiv` instructions. - [`div_round_half_up`](https://github.com/apache/datafusion-comet/blob/f9dd86d482d136426990186c2f1d42eef6a1ca59/native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs#L121-L146) rescales each `i256` product in `WideDecimalBinaryExpr`. It computes `value / divisor` and `value % divisor`, and each `i256` operator runs the long division in arrow-rs ([`div_rem_knuth`](https://github.com/apache/arrow-rs/blob/c7f60c612bd16ca5485cc9acddfd216c10ea9765/arrow-buffer/src/bigint/div.rs#L83)), so every row runs it twice. [`div_round_half_up_i128`](https://github.com/apache/datafusion-comet/blob/f9dd86d482d136426990186c2f1d42eef6a1ca59/native/spark-expr/src/conversion_funcs/string.rs#L521-L530) also uses `/` and `%`, but LLVM computes the remainder from the quotient there, and the same build makes one call to `__divti3`. ### Describe the potential solution In `pmod`, replace `(r + n) % n` with `r + n`. When `r` is negative it's in `(-n, 0)`, so `r + n` is already in `(0, n)` and can't overflow. The hashes in [`test_pmod`](https://github.com/apache/datafusion-comet/blob/f9dd86d482d136426990186c2f1d42eef6a1ca59/native/shuffle/src/comet_partitioning.rs#L135-L143) are all negative as `i32`, so the existing test covers that branch. In `div_round_half_up`, compute the remainder as `value.wrapping_sub(quot.wrapping_mul(divisor))`. `quot` truncates toward zero, so `|quot * divisor| <= |value|` and the product can't overflow. This uses only `i256` methods that the function already calls, so it doesn't need a new arrow-rs release. Once Comet uses an arrow-rs release with `i256::wrapping_div_rem` (apache/arrow-rs#11425), that method returns both results from one division. In a prototype on an Apple M5 Max, from apache/arrow-rs#11421: | Case | Today | One division | |---|---|---| | `pmod` on 8192 hashes, 13 partitions | 5.84 us | 4.59 us (1.3x) | | `pmod` on 8192 hashes, 16 partitions | 5.98 us | 4.52 us (1.3x) | | `wide_decimal` [`multiply`](https://github.com/apache/datafusion-comet/blob/8d0bb01a5a99924024e07c1722e6eddd9260e167/native/spark-expr/benches/wide_decimal.rs#L156), 30-digit products | 266.5 us | 233.8 us (1.1x) | | Multiply with 48-digit products | 398.1 us | 301.9 us (1.3x) | | Multiply with 42-digit products | 351.5 us | 234.5 us (1.5x) | The `pmod` numbers come from a standalone benchmark in the prototype, not from a Comet benchmark. The last two rows are cases the prototype added to the `wide_decimal` benchmark. ### Additional context apache/arrow-rs#11421 plans precomputed divisors that would replace both divisions. In the prototype they made `pmod` 4.0x faster and the multiply 2.6x to 4.2x faster. Comet could adopt them only after an arrow-rs release, and these changes don't depend on them. -- 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]
