fn criterion_benchmark()

in native/spark-expr/benches/aggregate.rs [41:113]

• 67 lines of code
• 6 McCabe index (conditional complexity)

fn criterion_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("aggregate");
    let num_rows = 8192;
    let batch = create_record_batch(num_rows);
    let mut batches = Vec::new();
    for _ in 0..10 {
        batches.push(batch.clone());
    }
    let partitions = &[batches];
    let c0: Arc<dyn PhysicalExpr> = Arc::new(Column::new("c0", 0));
    let c1: Arc<dyn PhysicalExpr> = Arc::new(Column::new("c1", 1));

    let rt = Runtime::new().unwrap();

    group.bench_function("avg_decimal_datafusion", |b| {
        let datafusion_sum_decimal = avg_udaf();
        b.to_async(&rt).iter(|| {
            black_box(agg_test(
                partitions,
                c0.clone(),
                c1.clone(),
                datafusion_sum_decimal.clone(),
                "avg",
            ))
        })
    });

    group.bench_function("avg_decimal_comet", |b| {
        let comet_avg_decimal = Arc::new(AggregateUDF::new_from_impl(AvgDecimal::new(
            DataType::Decimal128(38, 10),
            DataType::Decimal128(38, 10),
        )));
        b.to_async(&rt).iter(|| {
            black_box(agg_test(
                partitions,
                c0.clone(),
                c1.clone(),
                comet_avg_decimal.clone(),
                "avg",
            ))
        })
    });

    group.bench_function("sum_decimal_datafusion", |b| {
        let datafusion_sum_decimal = sum_udaf();
        b.to_async(&rt).iter(|| {
            black_box(agg_test(
                partitions,
                c0.clone(),
                c1.clone(),
                datafusion_sum_decimal.clone(),
                "sum",
            ))
        })
    });

    group.bench_function("sum_decimal_comet", |b| {
        let comet_sum_decimal = Arc::new(AggregateUDF::new_from_impl(
            SumDecimal::try_new(DataType::Decimal128(38, 10)).unwrap(),
        ));
        b.to_async(&rt).iter(|| {
            black_box(agg_test(
                partitions,
                c0.clone(),
                c1.clone(),
                comet_sum_decimal.clone(),
                "sum",
            ))
        })
    });

    group.finish();
}