sunchao commented on code in PR #5367:
URL: https://github.com/apache/datafusion-comet/pull/5367#discussion_r4016467234
##########
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:
[P2] Use the combined null count for binary pow dispatch
Could we base this dispatch on the combined null mask and reuse it for
evaluation? With independent null patterns, each operand can be below 75% null
while over 90% of output rows are null. This still selects `binary()`, which
computes powers for every masked payload.
For `pow(a + 2.5D, b + 2.5D)` with zero-filled input null slots, the
additions leave 2.5 under the preserved null bits. The new path evaluates those
expensive powers, while the base implementation skips them.
I compared the exact base/head power sources in optimized x86_64 Linux
pipelines with 8,192 rows, both additions inside the timed region, thin LTO,
one codegen unit, CPU affinity, and warm-up. Seven randomized
base/head/candidate triples per regression case were repeated with another
input seed:
| Nulls per operand | First run paired slowdown | Repeat |
| --- | ---: | ---: |
| About 70% | 42% | 42% |
| About 74% | 64% | 59% |
Every regression pair was slower on the head. Computing and reusing the
combined mask removed the slowdown. No-null and aligned-dense controls improved
on this PR. Raw zero-payload inputs also improved, and adding only the left
operand was approximately unchanged. The current benchmarks use aligned masks
or a non-null second operand and miss this independently nullable composition.
These are local expression-pipeline measurements, not whole-query
throughput. The harness used pinned Arrow 59.3 and DataFusion 55.0 because the
mirror lacks 55.1. The relevant DataFusion dispatch/type files are
byte-identical between those releases. Could we add this independent-mask
composed-expression benchmark alongside the fix?
--
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]
