fix: Filter out nulls properly in approx_percentile_cont_with_weight (#17780)

Author: Jefffrey

datafusion/functions-aggregate/src/approx_percentile_cont.rs

@@ -384,19 +384,23 @@ impl ApproxPercentileAccumulator {
         }
     }
 
-    // public for approx_percentile_cont_with_weight
+    // pub(crate) for approx_percentile_cont_with_weight
     pub(crate) fn max_size(&self) -> usize {
         self.digest.max_size()
     }
 
-    // public for approx_percentile_cont_with_weight
-    pub fn merge_digests(&mut self, digests: &[TDigest]) {
+    // pub(crate) for approx_percentile_cont_with_weight
+    pub(crate) fn merge_digests(&mut self, digests: &[TDigest]) {
         let digests = digests.iter().chain(std::iter::once(&self.digest));
         self.digest = TDigest::merge_digests(digests)
     }
 
-    // public for approx_percentile_cont_with_weight
-    pub fn convert_to_float(values: &ArrayRef) -> Result<Vec<f64>> {
+    // pub(crate) for approx_percentile_cont_with_weight
+    pub(crate) fn convert_to_float(values: &ArrayRef) -> Result<Vec<f64>> {
+        debug_assert!(
+            values.null_count() == 0,
+            "convert_to_float assumes nulls have already been filtered out"
+        );
         match values.data_type() {
             DataType::Float64 => {
                 let array = downcast_value!(values, Float64Array);
@@ -493,7 +497,7 @@ impl Accumulator for ApproxPercentileAccumulator {
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         // Remove any nulls before computing the percentile
         let mut values = Arc::clone(&values[0]);
-        if values.nulls().is_some() {
+        if values.null_count() > 0 {
             values = filter(&values, &is_not_null(&values)?)?;
         }
         let sorted_values = &arrow::compute::sort(&values, None)?;

datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs

@@ -21,6 +21,7 @@ use std::hash::Hash;
 use std::mem::size_of_val;
 use std::sync::Arc;
 
+use arrow::compute::{and, filter, is_not_null};
 use arrow::datatypes::FieldRef;
 use arrow::{array::ArrayRef, datatypes::DataType};
 use datafusion_common::ScalarValue;
@@ -268,15 +269,37 @@ impl Accumulator for ApproxPercentileWithWeightAccumulator {
     }
 
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
-        let means = &values[0];
-        let weights = &values[1];
+        let mut means = Arc::clone(&values[0]);
+        let mut weights = Arc::clone(&values[1]);
+        // If nulls are present in either array, need to filter those rows out in both arrays
+        match (means.null_count() > 0, weights.null_count() > 0) {
+            // Both have nulls
+            (true, true) => {
+                let predicate = and(&is_not_null(&means)?, &is_not_null(&weights)?)?;
+                means = filter(&means, &predicate)?;
+                weights = filter(&weights, &predicate)?;
+            }
+            // Only one has nulls
+            (false, true) => {
+                let predicate = &is_not_null(&weights)?;
+                means = filter(&means, predicate)?;
+                weights = filter(&weights, predicate)?;
+            }
+            (true, false) => {
+                let predicate = &is_not_null(&means)?;
+                means = filter(&means, predicate)?;
+                weights = filter(&weights, predicate)?;
+            }
+            // No nulls
+            (false, false) => {}
+        }
         debug_assert_eq!(
             means.len(),
             weights.len(),
             "invalid number of values in means and weights"
         );
-        let means_f64 = ApproxPercentileAccumulator::convert_to_float(means)?;
-        let weights_f64 = ApproxPercentileAccumulator::convert_to_float(weights)?;
+        let means_f64 = ApproxPercentileAccumulator::convert_to_float(&means)?;
+        let weights_f64 = ApproxPercentileAccumulator::convert_to_float(&weights)?;
         let mut digests: Vec<TDigest> = vec![];
         for (mean, weight) in means_f64.iter().zip(weights_f64.iter()) {
             digests.push(TDigest::new_with_centroid(

datafusion/sqllogictest/test_files/aggregate.slt

@@ -1540,6 +1540,19 @@ SELECT APPROX_PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY v) FROM (VALUES (CAST(
 ----
 NULL
 
+# percentile_cont_with_weight_with_nulls
+query I
+SELECT APPROX_PERCENTILE_CONT_WITH_WEIGHT(w, 0.5) WITHIN GROUP (ORDER BY v)
+FROM (VALUES (1, 1), (2, 1), (3, 1), (4, NULL), (NULL, 1), (NULL, NULL)) as t (v, w);
+----
+2
+
+# percentile_cont_with_weight_nulls_only
+query I
+SELECT APPROX_PERCENTILE_CONT_WITH_WEIGHT(1, 0.5) WITHIN GROUP (ORDER BY v) FROM (VALUES (CAST(NULL as INT))) as t (v);
+----
+NULL
+
 #
 # percentile_cont edge cases
 #