vortex/vortex-array/benches/expr/case_when_bench.rs at 8da5d6e8e838a8d37e133e5b0961e21b959e8871 · spiceai/vortex · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

#![allow(clippy::unwrap_used)]
#![allow(clippy::cast_possible_truncation)]

use std::sync::LazyLock;

use divan::Bencher;
use vortex_array::ArrayRef;
use vortex_array::Canonical;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::arrays::StructArray;
use vortex_array::expr::case_when;
use vortex_array::expr::get_item;
use vortex_array::expr::gt;
use vortex_array::expr::lit;
use vortex_array::expr::nested_case_when;
use vortex_array::expr::root;
use vortex_array::session::ArraySession;
use vortex_array::validity::Validity;
use vortex_buffer::Buffer;
use vortex_dtype::FieldNames;
use vortex_session::VortexSession;

static SESSION: LazyLock<VortexSession> =
    LazyLock::new(|| VortexSession::empty().with::<ArraySession>());

fn main() {
    divan::main();
}

fn make_struct_array(size: usize) -> ArrayRef {
    let data: Buffer<i32> = (0..size as i32).collect();
    let field = data.into_array();
    StructArray::try_new(
        FieldNames::from(["value"]),
        vec![field],
        size,
        Validity::NonNullable,
    )
    .unwrap()
    .into_array()
}

/// Benchmark a simple binary CASE WHEN with varying array sizes.
#[divan::bench(args = [1000, 10000, 100000])]
fn case_when_simple(bencher: Bencher, size: usize) {
    let array = make_struct_array(size);

    // CASE WHEN value > 500 THEN 100 ELSE 0 END
    let expr = case_when(
        gt(get_item("value", root()), lit(500i32)),
        lit(100i32),
        lit(0i32),
    );

    bencher
        .with_inputs(|| (&expr, &array))
        .bench_refs(|(expr, array)| {
            let mut ctx = SESSION.create_execution_ctx();
            array
                .apply(expr)
                .unwrap()
                .execute::<Canonical>(&mut ctx)
                .unwrap()
        });
}

/// Benchmark nested CASE WHEN with multiple conditions.
#[divan::bench(args = [1000, 10000, 100000])]
fn case_when_nested_3_conditions(bencher: Bencher, size: usize) {
    let array = make_struct_array(size);

    // CASE WHEN value > 750 THEN 3 WHEN value > 500 THEN 2 WHEN value > 250 THEN 1 ELSE 0 END
    let expr = nested_case_when(
        vec![
            (gt(get_item("value", root()), lit(750i32)), lit(3i32)),
            (gt(get_item("value", root()), lit(500i32)), lit(2i32)),
            (gt(get_item("value", root()), lit(250i32)), lit(1i32)),
        ],
        Some(lit(0i32)),
    );

    bencher
        .with_inputs(|| (&expr, &array))
        .bench_refs(|(expr, array)| {
            let mut ctx = SESSION.create_execution_ctx();
            array
                .apply(expr)
                .unwrap()
                .execute::<Canonical>(&mut ctx)
                .unwrap()
        });
}

/// Benchmark CASE WHEN where all conditions are true (short-circuit path).
#[divan::bench(args = [1000, 10000, 100000])]
fn case_when_all_true(bencher: Bencher, size: usize) {
    let array = make_struct_array(size);

    // CASE WHEN value >= 0 THEN 100 ELSE 0 END (always true for our data)
    let expr = case_when(
        gt(get_item("value", root()), lit(-1i32)),
        lit(100i32),
        lit(0i32),
    );

    bencher
        .with_inputs(|| (&expr, &array))
        .bench_refs(|(expr, array)| {
            let mut ctx = SESSION.create_execution_ctx();
            array
                .apply(expr)
                .unwrap()
                .execute::<Canonical>(&mut ctx)
                .unwrap()
        });
}

/// Benchmark CASE WHEN where all conditions are false (short-circuit path).
#[divan::bench(args = [1000, 10000, 100000])]
fn case_when_all_false(bencher: Bencher, size: usize) {
    let array = make_struct_array(size);

    // CASE WHEN value > 1000000 THEN 100 ELSE 0 END (always false for our data)
    let expr = case_when(
        gt(get_item("value", root()), lit(1_000_000i32)),
        lit(100i32),
        lit(0i32),
    );

    bencher
        .with_inputs(|| (&expr, &array))
        .bench_refs(|(expr, array)| {
            let mut ctx = SESSION.create_execution_ctx();
            array
                .apply(expr)
                .unwrap()
                .execute::<Canonical>(&mut ctx)
                .unwrap()
        });
}