Improve speed of first_value by implementing special GroupsAccumulator

Author: UBarney

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -15,11 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::str;
 use std::sync::Arc;
 
 use crate::fuzz_cases::aggregation_fuzzer::{
-    AggregationFuzzerBuilder, ColumnDescr, DatasetGeneratorConfig, QueryBuilder,
+    AggregationFuzzerBuilder, ColumnDescr,
+    DatasetGeneratorConfig, QueryBuilder,
 };
 
 use arrow::array::{types::Int64Type, Array, ArrayRef, AsArray, Int64Array, RecordBatch};
@@ -88,6 +88,141 @@ async fn test_min() {
         .await;
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_first_val() {
+    let mut data_gen_config: DatasetGeneratorConfig = baseline_config();
+
+    // for ele in data_gen_config.columns.iter_mut() {
+    //     if ele.get_max_num_distinct().is_none() {
+    //         ele.with_max_num_distinct(usize::MAX);
+    //     }
+    // }
+
+    for i in 0..data_gen_config.columns.len() {
+        // data_gen_config.columns[i]
+        if data_gen_config.columns[i].get_max_num_distinct().is_none() {
+            data_gen_config.columns[i] = data_gen_config.columns[i]
+                .clone()
+                .with_max_num_distinct(usize::MAX);
+        }
+    }
+
+    // data_gen_config.columns.iter_mut().for_each(|ele| {
+    //     if ele.get_max_num_distinct().is_none() {
+    //         ele.with_max_num_distinct(usize::MAX);
+    //     }
+    // });
+
+    // Queries like SELECT max(a) FROM fuzz_table GROUP BY b
+    let query_builder = QueryBuilder::new()
+        .with_table_name("fuzz_table")
+        .with_aggregate_function("first_value")
+        // max works on all column types
+        .with_aggregate_arguments(data_gen_config.all_columns())
+        .set_group_by_columns(
+            data_gen_config.all_columns(), // .into_iter()
+                                           // .filter(|x| !x.contains("u8")),
+        );
+
+    AggregationFuzzerBuilder::from(data_gen_config)
+        .add_query_builder(query_builder)
+        .build()
+        .run()
+        .await;
+}
+
+// #[tokio::test(flavor = "multi_thread")]
+// async fn test_bad_case() {
+//     let sql="SELECT u8_low, first_value(i16 order by u8_low DESC, float64 DESC, large_binary DESC,utf8_low DESC,interval_year_month ASC) RESPECT NULLS as col2 FROM parquet_table GROUP BY u8_low
+
+// ";
+
+//     let (baseline_ctx, random_ctx) =
+//         gen_ctx(baseline_config(), "/tmp/SYuFrB4.parquet").await;
+
+//     //     {let explain = "explain ".to_owned() + sql;
+//     //     let baseline_plan = run_sql(&explain, &baseline_ctx).await.unwrap();
+//     //     let plan = run_sql(&explain, &random_ctx).await.unwrap();
+//     //     println!("{}", pretty_format_batches(&baseline_plan).unwrap());
+//     //     println!("{}", pretty_format_batches(&plan).unwrap());
+//     // }
+//     let baseline_res = run_sql(sql, &baseline_ctx).await.unwrap();
+
+//     println!("<==========================================================>");
+//     println!("<==========================================================>");
+//     println!("<==========================================================>");
+//     println!("<==========================================================>");
+
+//     let got = run_sql(sql, &random_ctx).await.unwrap();
+
+//     // println!("{}", pretty_format_batches(&baseline_res).unwrap());
+
+//     println!("\n*********************\n");
+
+//     let baseline_count = baseline_res.iter().map(|x| x.num_rows()).sum::<usize>();
+
+//     let got_count = got.iter().map(|x| x.num_rows()).sum::<usize>();
+
+//     if baseline_count != got_count {
+//         panic!(
+//             "baseline_res.len()!=got.len() {} {}",
+//             baseline_count, got_count
+//         );
+//     }
+
+//     check_equality_of_batches(&baseline_res, &got).unwrap();
+// }
+
+// #[tokio::test]
+// async fn test_dev_first() -> Result<()> {
+//     let ctx = SessionContext::new();
+
+//     ctx.sql(
+//         "CREATE TABLE array_agg_order_list_table AS VALUES
+//   ('w', 2, [1,2,3], 10),
+//   ('w', 1, [9,5,2], 20),
+//   ('w', 1, [3,2,5], 30),
+//   ('b', 2, [4,5,6], 20),
+//   ('b', 1, [7,8,9], 30)
+// ;
+// ",
+//     )
+//     .await
+//     .unwrap();
+
+//     ctx.sql("select column1, first_value(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1;").and_then(|x|async move{x.show().await}).await.map_err(|x|{
+//         println!("");
+//         eprintln!("{x}");
+//         x
+//     })?;
+
+//     // RUST_BACKTRACE=1 cargo test --features=backtrace  --package datafusion --test fuzz -- fuzz_cases::aggregate_fuzz::test_first_val  --exact  --nocapture
+
+//     Ok(())
+// }
+
+#[tokio::test]
+async fn test_get_backtrace_for_failed_code() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    let sql = "select row_numer() over (partition by a order by a) from (select 1 a);";
+
+    match ctx.sql(sql).await {
+        Ok(result) => result.show().await?,
+        Err(e) => {
+            eprintln!("{e}");
+        }
+    };
+
+    // $ RUST_BACKTRACE=1 cargo test --features=backtrace --package datafusion --lib -- fuzz_cases::aggregate_fuzz::test_get_backtrace_for_failed_code --exact --nocapture
+
+    // RUST_BACKTRACE=1 cargo test --features=backtrace  --package datafusion --test fuzz -- fuzz_cases::aggregate_fuzz::test_dev_first  --exact  --nocapture
+
+    // RUST_BACKTRACE=1 cargo test --features=backtrace  --package datafusion --test fuzz -- fuzz_cases::aggregate_fuzz::test_get_backtrace_for_failed_code  --exact  --nocapture
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_max() {
     let data_gen_config = baseline_config();
@@ -216,27 +351,25 @@ fn baseline_config() -> DatasetGeneratorConfig {
             "interval_day_time",
             DataType::Interval(IntervalUnit::DayTime),
         ),
-        ColumnDescr::new(
-            "interval_month_day_nano",
-            DataType::Interval(IntervalUnit::MonthDayNano),
-        ),
+        // ColumnDescr::new(
+        //     "interval_month_day_nano",
+        //     DataType::Interval(IntervalUnit::MonthDayNano),
+        // ),
         // begin decimal columns
         ColumnDescr::new("decimal128", {
             // Generate valid precision and scale for Decimal128 randomly.
             let precision: u8 = rng.gen_range(1..=DECIMAL128_MAX_PRECISION);
             // It's safe to cast `precision` to i8 type directly.
-            let scale: i8 = rng.gen_range(
-                i8::MIN..=std::cmp::min(precision as i8, DECIMAL128_MAX_SCALE),
-            );
+            let scale: i8 =
+                rng.gen_range(0..=std::cmp::min(precision as i8, DECIMAL128_MAX_SCALE));
             DataType::Decimal128(precision, scale)
         }),
         ColumnDescr::new("decimal256", {
             // Generate valid precision and scale for Decimal256 randomly.
             let precision: u8 = rng.gen_range(1..=DECIMAL256_MAX_PRECISION);
             // It's safe to cast `precision` to i8 type directly.
-            let scale: i8 = rng.gen_range(
-                i8::MIN..=std::cmp::min(precision as i8, DECIMAL256_MAX_SCALE),
-            );
+            let scale: i8 =
+                rng.gen_range(0..=std::cmp::min(precision as i8, DECIMAL256_MAX_SCALE));
             DataType::Decimal256(precision, scale)
         }),
         // begin string columns

datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs

@@ -228,6 +228,10 @@ impl ColumnDescr {
         }
     }
 
+    pub fn get_max_num_distinct(&self) -> Option<usize> {
+        self.max_num_distinct
+    }
+
     /// set the maximum number of distinct values in this column
     ///
     /// If `None`, the number of distinct values is randomly selected between 1