Merge branch 'main' into aggregation-fuzzer-spill-coverage

Author: naman-modi

datafusion/functions-aggregate/Cargo.toml

@@ -100,5 +100,9 @@ harness = false
 name = "sliding_max"
 harness = false
 
+[[bench]]
+name = "variance"
+harness = false
+
 [features]
 force_hash_collisions = ["datafusion-common/force_hash_collisions"]

datafusion/functions-aggregate/benches/variance.rs

@@ -0,0 +1,83 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Float64Array};
+use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
+use datafusion_expr::Accumulator;
+use datafusion_functions_aggregate::variance::VarianceAccumulator;
+use datafusion_functions_aggregate_common::stats::StatsType;
+
+const BATCH_SIZE: usize = 8192;
+
+fn batch_array(null_stride: Option<usize>) -> ArrayRef {
+    let values = (0..BATCH_SIZE)
+        .map(|idx| {
+            if null_stride.is_some_and(|stride| idx % stride == 0) {
+                None
+            } else {
+                Some(idx as f64)
+            }
+        })
+        .collect::<Vec<_>>();
+    Arc::new(Float64Array::from(values)) as ArrayRef
+}
+
+fn update_bench(c: &mut Criterion, name: &str, batch: &ArrayRef) {
+    c.bench_function(name, |b| {
+        b.iter(|| {
+            let mut acc = VarianceAccumulator::try_new(StatsType::Sample).unwrap();
+            acc.update_batch(std::slice::from_ref(batch)).unwrap();
+            black_box(acc.evaluate().unwrap())
+        })
+    });
+}
+
+fn retract_bench(c: &mut Criterion, name: &str, batch: &ArrayRef) {
+    c.bench_function(name, |b| {
+        b.iter_batched(
+            || {
+                let mut acc = VarianceAccumulator::try_new(StatsType::Sample).unwrap();
+                // Accumulate two batches so that retracting one leaves the
+                // accumulator with rows remaining, as in a sliding window.
+                acc.update_batch(std::slice::from_ref(batch)).unwrap();
+                acc.update_batch(std::slice::from_ref(batch)).unwrap();
+                acc
+            },
+            |mut acc| {
+                acc.retract_batch(std::slice::from_ref(batch)).unwrap();
+                black_box(acc.evaluate().unwrap())
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+fn variance_benchmark(c: &mut Criterion) {
+    let no_nulls = batch_array(None);
+    let with_nulls = batch_array(Some(10));
+
+    update_bench(c, "variance update_batch f64 no_nulls", &no_nulls);
+    update_bench(c, "variance update_batch f64 with_nulls", &with_nulls);
+    retract_bench(c, "variance retract_batch f64 no_nulls", &no_nulls);
+    retract_bench(c, "variance retract_batch f64 with_nulls", &with_nulls);
+}
+
+criterion_group!(benches, variance_benchmark);
+criterion_main!(benches);

datafusion/functions-aggregate/src/variance.rs

@@ -326,6 +326,23 @@ fn update(count: u64, mean: f64, m2: f64, value: f64) -> (u64, f64, f64) {
     (new_count, new_mean, new_m2)
 }
 
+/// Inverse of [`update`]: removes a previously accumulated value. Retracting
+/// from a state with one or zero values resets the state to empty.
+#[inline]
+fn retract(count: u64, mean: f64, m2: f64, value: f64) -> (u64, f64, f64) {
+    if count <= 1 {
+        return (0, 0.0, 0.0);
+    }
+
+    let new_count = count - 1;
+    let delta1 = mean - value;
+    let new_mean = delta1 / new_count as f64 + mean;
+    let delta2 = new_mean - value;
+    let new_m2 = m2 - delta1 * delta2;
+
+    (new_count, new_mean, new_m2)
+}
+
 impl Accumulator for VarianceAccumulator {
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
         Ok(vec![
@@ -348,22 +365,8 @@ impl Accumulator for VarianceAccumulator {
     fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         let arr = as_float64_array(&values[0])?;
         for value in arr.iter().flatten() {
-            if self.count <= 1 {
-                self.count = 0;
-                self.mean = 0.0;
-                self.m2 = 0.0;
-                continue;
-            }
-
-            let new_count = self.count - 1;
-            let delta1 = self.mean - value;
-            let new_mean = delta1 / new_count as f64 + self.mean;
-            let delta2 = new_mean - value;
-            let new_m2 = self.m2 - delta1 * delta2;
-
-            self.count -= 1;
-            self.mean = new_mean;
-            self.m2 = new_m2;
+            (self.count, self.mean, self.m2) =
+                retract(self.count, self.mean, self.m2, value)
         }
 
         Ok(())
@@ -691,6 +694,75 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn update_batch_ignores_nulls() -> Result<()> {
+        // An array with nulls must accumulate the same values as a dense
+        // array of its non-null values.
+        let dense: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]));
+        let sparse: ArrayRef = Arc::new(Float64Array::from(vec![
+            Some(1.0),
+            None,
+            Some(2.0),
+            Some(3.0),
+            None,
+            Some(4.0),
+        ]));
+
+        let mut dense_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        dense_acc.update_batch(std::slice::from_ref(&dense))?;
+        let mut sparse_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        sparse_acc.update_batch(std::slice::from_ref(&sparse))?;
+
+        // Sample variance of {1, 2, 3, 4} is 5/3 (all steps are exact in f64).
+        assert_eq!(dense_acc.evaluate()?, ScalarValue::Float64(Some(5.0 / 3.0)));
+        assert_eq!(dense_acc.evaluate()?, sparse_acc.evaluate()?);
+        Ok(())
+    }
+
+    #[test]
+    fn retract_batch_ignores_nulls() -> Result<()> {
+        let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]));
+        let dense_retract: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0]));
+        let sparse_retract: ArrayRef =
+            Arc::new(Float64Array::from(vec![Some(1.0), None, Some(2.0)]));
+
+        let mut dense_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        dense_acc.update_batch(std::slice::from_ref(&values))?;
+        dense_acc.retract_batch(std::slice::from_ref(&dense_retract))?;
+        let mut sparse_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        sparse_acc.update_batch(std::slice::from_ref(&values))?;
+        sparse_acc.retract_batch(std::slice::from_ref(&sparse_retract))?;
+
+        // Sample variance of the remaining {3, 4} is 0.5 (all steps are exact
+        // in f64).
+        assert_eq!(dense_acc.evaluate()?, ScalarValue::Float64(Some(0.5)));
+        assert_eq!(dense_acc.evaluate()?, sparse_acc.evaluate()?);
+        Ok(())
+    }
+
+    #[test]
+    fn retract_batch_resets_when_underflowing() -> Result<()> {
+        // Retracting more values than were accumulated resets to the empty
+        // state, with or without nulls in the retracted batch.
+        let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0]));
+        let dense_retract: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0]));
+        let sparse_retract: ArrayRef = Arc::new(Float64Array::from(vec![
+            Some(1.0),
+            None,
+            Some(2.0),
+            Some(3.0),
+        ]));
+
+        for retract in [&dense_retract, &sparse_retract] {
+            let mut acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+            acc.update_batch(std::slice::from_ref(&values))?;
+            acc.retract_batch(std::slice::from_ref(retract))?;
+            assert_eq!(acc.get_count(), 0);
+            assert_eq!(acc.evaluate()?, ScalarValue::Float64(None));
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_groups_accumulator_merge_empty_states() -> Result<()> {
         let state_1 = vec![

datafusion/functions/src/datetime/to_date.rs

@@ -61,7 +61,7 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo
         name = "format_n",
         description = r"Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order
   they appear with the first successful one being returned. If none of the formats successfully parse the expression
-  an error will be returned."
+  an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL."
     )
 )]
 #[derive(Debug, PartialEq, Eq, Hash)]
@@ -519,4 +519,44 @@ mod tests {
             panic!("Conversion of {date_str} succeeded, but should have failed. ");
         }
     }
+
+    /// A NULL format must be skipped even when its slot still holds parseable
+    /// bytes, otherwise it can silently win over a later valid format.
+    #[test]
+    fn test_to_date_null_format_slot_retaining_bytes() {
+        use arrow::buffer::NullBuffer;
+
+        // The first format physically holds "%d/%m/%Y", but is marked NULL.
+        let (offsets, values, _) =
+            GenericStringArray::<i32>::from(vec!["%d/%m/%Y"]).into_parts();
+        let formats =
+            GenericStringArray::new(offsets, values, Some(NullBuffer::new_null(1)));
+        assert!(formats.is_null(0));
+        assert_eq!(formats.value(0), "%d/%m/%Y");
+
+        // Without the validity check, the first format parses this as 2023-02-01
+        // and incorrectly wins over the valid second format.
+        let values = GenericStringArray::<i32>::from(vec!["01/02/2023"]);
+        let fallback_formats = GenericStringArray::<i32>::from(vec!["%m/%d/%Y"]);
+        let res = invoke_to_date_with_args(
+            vec![
+                ColumnarValue::Array(Arc::new(values)),
+                ColumnarValue::Array(Arc::new(formats)),
+                ColumnarValue::Array(Arc::new(fallback_formats)),
+            ],
+            1,
+        )
+        .unwrap();
+
+        let ColumnarValue::Array(res) = res else {
+            panic!("expected an array result");
+        };
+        let res = res.as_any().downcast_ref::<Date32Array>().unwrap();
+
+        assert!(!res.is_null(0));
+        assert_eq!(
+            res.value(0),
+            Date32Type::parse_formatted("01/02/2023", "%m/%d/%Y").unwrap()
+        );
+    }
 }

datafusion/functions/src/datetime/to_timestamp.rs

@@ -81,7 +81,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo
         description = r#"
 Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression.
 Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully
-parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
+parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL.
+Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
 only supported at the end of the string preceded by a space.
 "#
     )
@@ -131,7 +132,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo
         description = r#"
 Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression.
 Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully
-parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
+parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL.
+Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
 only supported at the end of the string preceded by a space.
 "#
     )
@@ -181,7 +183,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo
         description = r#"
 Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression.
 Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully
-parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
+parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL.
+Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
 only supported at the end of the string preceded by a space.
 "#
     )
@@ -231,7 +234,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo
         description = r#"
 Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression.
 Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully
-parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
+parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL.
+Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
 only supported at the end of the string preceded by a space.
 "#
     )
@@ -280,7 +284,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo
         description = r#"
 Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression.
 Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully
-parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
+parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL.
+Note: parsing of named timezones (e.g. 'America/New_York') using %Z is
 only supported at the end of the string preceded by a space.
 "#
     )

datafusion/functions/src/datetime/to_unixtime.rs

@@ -56,7 +56,7 @@ Integers, unsigned integers, and floats are interpreted as seconds since the uni
     ),
     argument(
         name = "format_n",
-        description = "Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression an error will be returned."
+        description = "Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL."
     )
 )]
 #[derive(Debug, PartialEq, Eq, Hash)]

datafusion/sqllogictest/test_files/datetime/dates.slt

@@ -319,6 +319,14 @@ ORDER BY id
 1 2020-09-08
 2 NULL
 
+# Skipping NULL formats does not mask a later parse error.
+query error DataFusion error: Execution error: Error parsing timestamp from '2020\-09\-08' using format '%q': trailing input
+SELECT to_date('2020-09-08', NULL::VARCHAR, '%q')
+
+# Invalid format types are rejected before NULL input propagation.
+query error DataFusion error: Execution error: to_date function unsupported data type at index 1: Int64
+SELECT to_date(NULL::VARCHAR, 12345)
+
 statement ok
 create table ts_utf8_data(ts varchar(100), format varchar(100)) as values
   ('2020-09-08 12/00/00+00:00', '%Y-%m-%d %H/%M/%S%#z'),