sunchao commented on code in PR #5367:
URL: https://github.com/apache/datafusion-comet/pull/5367#discussion_r3836656343
##########
native/spark-expr/src/math_funcs/pow.rs:
##########
@@ -42,86 +45,55 @@ 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 uses
[`unary`] so the
+/// scalar is not broadcast. A null scalar short-circuits to an all-null array.
+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)))
- }
- (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)))
- }
- (ColumnarValue::Array(base_arr), ColumnarValue::Scalar(exp_scalar)) =>
{
- let bases = as_f64_array(base_arr)?;
- let result: Float64Array = match as_f64_scalar(exp_scalar)? {
- Some(exp) => bases
- .iter()
- .map(|b| b.map(|base| spark_powf(base, exp)))
- .collect(),
- None => Float64Array::new_null(base_arr.len()),
- };
- 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 {
+ unary(right, |exp| spark_powf(left.value(0), exp))
+ }
}
- (ColumnarValue::Scalar(base_scalar),
ColumnarValue::Scalar(exp_scalar)) => {
- let result = match (as_f64_scalar(base_scalar)?,
as_f64_scalar(exp_scalar)?) {
- (Some(base), Some(exp)) =>
ScalarValue::Float64(Some(spark_powf(base, exp))),
- _ => ScalarValue::Float64(None),
- };
- Ok(ColumnarValue::Scalar(result))
+ (false, true) => {
+ if right.is_null(0) {
+ Float64Array::new_null(left.len())
+ } else {
+ unary(left, |base| spark_powf(base, right.value(0)))
+ }
}
- }
+ _ => binary(left, right, spark_powf)?,
Review Comment:
[P2] Keep null-skipping for dense nullable inputs
Could we retain a null-aware path for dense inputs and benchmark a nullable
child expression? For `pow(a + 2.5D, b)`, with nullable double `a` and finite
fractional exponents in `b`, Arrow addition preserves the null bits but changes
their underlying values from zero to `2.5`. These `unary`/`binary` calls then
compute powers for null slots that the previous iterator skipped.
I compared the exact base/head `pow.rs` files in a standalone optimized
harness with Arrow 58.4 and DataFusion 54.1, using Comet's release optimization
settings (thin LTO and one codegen unit). On arm64 macOS, with 8,192-row
batches, medians from nine alternating rounds after warm-up showed the
following, including the addition in the timed pipeline:
| Null fraction | Base | This head | Slowdown |
| --- | ---: | ---: | ---: |
| About 90% | 26.0 us | 47.9 us | 1.84x |
| About 99% | 23.8 us | 47.0 us | 1.97x |
Separate repeats and an independent rerun reproduced the regression, while
the no-null pipeline improved. Nullness and valid result bits matched. These
are local kernel/pipeline measurements, not full Spark-query or
production-workload measurements.
The current benchmarks construct null slots directly from `None`, leaving
zero payloads, and stop at 50% nulls. That misses this composed-expression
case. Could we cover it and preserve null-skipping where evaluating those
masked values is more expensive than the iterator overhead?
--
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]
