0lai0 commented on code in PR #5367:
URL: https://github.com/apache/datafusion-comet/pull/5367#discussion_r4039101971
##########
native/spark-expr/src/math_funcs/pow.rs:
##########
@@ -42,86 +60,131 @@ fn spark_powf(base: f64, exp: f64) -> f64 {
/// Unlike DataFusion's `power`, `pow(0, -1)` returns `Infinity` rather than
erroring. Only null
/// inputs produce null; otherwise every result is the `spark_powf` value.
pub fn spark_pow(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
- if args.len() != 2 {
- return Err(DataFusionError::Internal(format!(
- "spark_pow requires 2 arguments, got {}",
- args.len()
- )));
- }
+ let [base, exp] = take_function_args("spark_pow", args)?;
+ apply(base, exp, spark_pow_kernel)
+}
- fn as_f64_array(
- value: &Arc<dyn arrow::array::Array>,
- ) -> Result<&Float64Array, DataFusionError> {
- value
- .as_any()
- .downcast_ref::<Float64Array>()
- .ok_or_else(|| {
- DataFusionError::Internal(format!(
- "spark_pow expected Float64, got {:?}",
- value.data_type()
- ))
- })
- }
+fn as_f64_array(array: &dyn Array) -> Result<&Float64Array, ArrowError> {
+ array
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .ok_or_else(|| {
+ ArrowError::ComputeError(format!(
+ "spark_pow expected Float64, got {:?}",
+ array.data_type()
+ ))
+ })
+}
- fn as_f64_scalar(scalar: &ScalarValue) -> Result<Option<f64>,
DataFusionError> {
- match scalar {
- ScalarValue::Float64(v) => Ok(*v),
- _ => Err(DataFusionError::Internal(format!(
- "spark_pow expected Float64 scalar, got {scalar:?}",
- ))),
- }
- }
+/// Array/array uses [`binary`] over [`spark_powf`]. Scalar/array and
array/scalar use
+/// [`unary`] so the scalar is not broadcast. A null scalar on either side
short-circuits
+/// to an all-null array. When null density exceeds the threshold (see
[`is_dense_null`])
+/// we skip masked slots to avoid running `spark_powf` on carried-over payload
from an
+/// upstream Arrow op.
+fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef,
ArrowError> {
+ let (left, left_is_scalar) = lhs.get();
+ let (right, right_is_scalar) = rhs.get();
+ let left = as_f64_array(left)?;
+ let right = as_f64_array(right)?;
- match (&args[0], &args[1]) {
- (ColumnarValue::Array(base_arr), ColumnarValue::Array(exp_arr)) => {
- let bases = as_f64_array(base_arr)?;
- let exps = as_f64_array(exp_arr)?;
- let result: Float64Array = bases
- .iter()
- .zip(exps.iter())
- .map(|(b, e)| match (b, e) {
- (Some(base), Some(exp)) => Some(spark_powf(base, exp)),
- _ => None,
- })
- .collect();
- Ok(ColumnarValue::Array(Arc::new(result)))
+ let result = match (left_is_scalar, right_is_scalar) {
+ (true, false) => {
+ if left.is_null(0) {
+ Float64Array::new_null(right.len())
+ } else if is_dense_null(right.null_count(), right.len()) {
+ pow_scalar_array_null_aware(left.value(0), right)
+ } else {
+ unary(right, |exp| spark_powf(left.value(0), exp))
+ }
+ }
+ (false, true) => {
+ if right.is_null(0) {
+ Float64Array::new_null(left.len())
+ } else if is_dense_null(left.null_count(), left.len()) {
+ pow_array_scalar_null_aware(left, right.value(0))
+ } else {
+ unary(left, |base| spark_powf(base, right.value(0)))
+ }
}
- (ColumnarValue::Scalar(base_scalar), ColumnarValue::Array(exp_arr)) =>
{
- let exps = as_f64_array(exp_arr)?;
- let result: Float64Array = match as_f64_scalar(base_scalar)? {
- Some(base) => exps
- .iter()
- .map(|e| e.map(|exp| spark_powf(base, exp)))
- .collect(),
- None => Float64Array::new_null(exp_arr.len()),
- };
- Ok(ColumnarValue::Array(Arc::new(result)))
+ _ => {
+ // Effective null count of the output is bounded below by
max(left, right).
+ // Using the max avoids a full NullBuffer::union scan just for the
density
+ // check; the union still happens if we actually take the
null-aware path.
+ let approx_nulls = left.null_count().max(right.null_count());
Review Comment:
Fixed in b9197b34e and 5be638d03.
- `pow_binary` builds `NullBuffer::union` of both masks once and reuses it
for the result, so dispatch sees rows that are null only on the other side.
- While checking this against base, I found the 75% threshold itself was
wrong: between roughly 50% and 75% nulls with a real payload in null slots,
head was still up to 30% slower than base. A sweep of evaluate-every-slot vs
skip-null-slots (random masks, 5%–75% nulls, zero-filled and payload-carrying
null slots, array/array and scalar/array) showed skipping matched or won at
every density from 5% upward, since a `spark_powf` call costs far more than
walking the bitmap. So the threshold is gone: null slots are skipped whenever
the output has nulls.
- A null `Float64` scalar now short-circuits before `apply()`, which was
materializing the scalar as a one-element array first (+12% on that path).
Added `spark_pow: pipeline pow(a + 2.5D, b + 2.5D) independent nulls
{10,30,50,70}%` (both additions timed), extended the other null sweeps down to
10%, and added unit tests for independent masks (70%/70%, 91% null output) and
an all-valid null buffer.
Base (main) vs head, interleaved, 3 rounds, 8,192 rows, local macOS aarch64,
expression pipeline only:
| Case | Base | Head | Change |
|---|---:|---:|---:|
| `pow(a + 2.5D, b + 2.5D)` independent nulls 10% | 52.4µs | 35.5µs | −32% |
| `pow(a + 2.5D, b + 2.5D)` independent nulls 30% | 40.9µs | 22.9µs | −44% |
| `pow(a + 2.5D, b + 2.5D)` independent nulls 50% | 32.6µs | 13.5µs | −59% |
| `pow(a + 2.5D, b + 2.5D)` independent nulls 70% | 27.6µs | 8.0µs | −71% |
| `pow(a + 2.5D, b)` nulls 50% | 36.3µs | 20.3µs | −44% |
| `pow(a + 2.5D, b)` nulls 70% | 29.4µs | 13.3µs | −55% |
| array/array no nulls | 40.5µs | 29.0µs | −28% |
| scalar/array no nulls | 43.2µs | 33.8µs | −22% |
| array/array dense nulls (aligned) | 33.5µs | 14.7µs | −56% |
| null scalar short-circuit (5 rounds) | 1.32µs | 1.36µs | +3.7% |
All 44 `spark_pow` benchmarks are faster than base except the null-scalar
short-circuit, which is within ~50ns.
--
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]
