|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! Benchmarks for the `list_sum` scalar function over `List`, `ListView`, and |
| 5 | +//! `FixedSizeList` inputs. |
| 6 | +
|
| 7 | +#![expect(clippy::unwrap_used)] |
| 8 | +#![expect(clippy::cast_possible_truncation)] |
| 9 | + |
| 10 | +use std::sync::LazyLock; |
| 11 | + |
| 12 | +use divan::Bencher; |
| 13 | +use rand::RngExt; |
| 14 | +use rand::SeedableRng; |
| 15 | +use rand::distr::Uniform; |
| 16 | +use rand::rngs::StdRng; |
| 17 | +use vortex_array::ArrayRef; |
| 18 | +use vortex_array::Canonical; |
| 19 | +use vortex_array::IntoArray; |
| 20 | +use vortex_array::VortexSessionExecute; |
| 21 | +use vortex_array::arrays::BoolArray; |
| 22 | +use vortex_array::arrays::FixedSizeListArray; |
| 23 | +use vortex_array::arrays::ListArray; |
| 24 | +use vortex_array::arrays::ListViewArray; |
| 25 | +use vortex_array::arrays::PrimitiveArray; |
| 26 | +use vortex_array::expr::list_sum; |
| 27 | +use vortex_array::expr::root; |
| 28 | +use vortex_array::validity::Validity; |
| 29 | +use vortex_buffer::Buffer; |
| 30 | +use vortex_session::VortexSession; |
| 31 | + |
| 32 | +fn main() { |
| 33 | + divan::main(); |
| 34 | +} |
| 35 | + |
| 36 | +static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session); |
| 37 | + |
| 38 | +const BASE_LIST_SIZE: usize = 8; |
| 39 | + |
| 40 | +const SMALL: usize = 100; |
| 41 | +const MEDIUM: usize = 10_000; |
| 42 | +const LARGE: usize = 1_000_000; |
| 43 | + |
| 44 | +/// A uniformly-random partition of `num_lists * BASE_LIST_SIZE` elements into `num_lists` |
| 45 | +/// lists, plus a validity mask with ~1/8 of lists null at random positions. |
| 46 | +fn random_lists(num_lists: usize) -> (Vec<i32>, Validity) { |
| 47 | + let mut rng = StdRng::seed_from_u64(num_lists as u64); |
| 48 | + let total = (num_lists * BASE_LIST_SIZE) as i32; |
| 49 | + |
| 50 | + let cut_dist = Uniform::new_inclusive(0i32, total).unwrap(); |
| 51 | + let mut cuts: Vec<i32> = (0..num_lists - 1).map(|_| rng.sample(cut_dist)).collect(); |
| 52 | + cuts.sort_unstable(); |
| 53 | + let mut sizes = Vec::with_capacity(num_lists); |
| 54 | + let mut prev = 0i32; |
| 55 | + for cut in cuts { |
| 56 | + sizes.push(cut - prev); |
| 57 | + prev = cut; |
| 58 | + } |
| 59 | + sizes.push(total - prev); |
| 60 | + |
| 61 | + let null_dist = Uniform::new(0u32, 8).unwrap(); |
| 62 | + let valid = (0..num_lists).map(|_| rng.sample(null_dist) != 0); |
| 63 | + ( |
| 64 | + sizes, |
| 65 | + Validity::Array(BoolArray::from_iter(valid).into_array()), |
| 66 | + ) |
| 67 | +} |
| 68 | + |
| 69 | +/// `0..total` as `i32` elements, with ~1/8 nulled at random positions when `nullable`. |
| 70 | +fn make_elements(total: i32, nullable: bool) -> ArrayRef { |
| 71 | + if !nullable { |
| 72 | + return PrimitiveArray::from_iter(0..total).into_array(); |
| 73 | + } |
| 74 | + let mut rng = StdRng::seed_from_u64(total as u64); |
| 75 | + let null_dist = Uniform::new(0u32, 8).unwrap(); |
| 76 | + PrimitiveArray::from_option_iter((0..total).map(|v| (rng.sample(null_dist) != 0).then_some(v))) |
| 77 | + .into_array() |
| 78 | +} |
| 79 | + |
| 80 | +/// A canonical `List<i32>` of `num_lists` variable-length lists, ~1/8 of them null. |
| 81 | +fn make_list(num_lists: usize, nullable_elements: bool) -> ArrayRef { |
| 82 | + let (sizes, validity) = random_lists(num_lists); |
| 83 | + let total: i32 = sizes.iter().sum(); |
| 84 | + let elements = make_elements(total, nullable_elements); |
| 85 | + let offsets: Buffer<i32> = std::iter::once(0) |
| 86 | + .chain(sizes.iter().scan(0i32, |acc, &s| { |
| 87 | + *acc += s; |
| 88 | + Some(*acc) |
| 89 | + })) |
| 90 | + .collect(); |
| 91 | + ListArray::try_new(elements, offsets.into_array(), validity) |
| 92 | + .unwrap() |
| 93 | + .into_array() |
| 94 | +} |
| 95 | + |
| 96 | +/// A gapless `ListView<i32>` of `num_lists` variable-length lists, ~1/8 of them null. |
| 97 | +fn make_listview(num_lists: usize) -> ArrayRef { |
| 98 | + let (sizes, validity) = random_lists(num_lists); |
| 99 | + let total: i32 = sizes.iter().sum(); |
| 100 | + let elements = make_elements(total, false); |
| 101 | + let offsets: Buffer<i32> = sizes |
| 102 | + .iter() |
| 103 | + .scan(0i32, |acc, &s| { |
| 104 | + let start = *acc; |
| 105 | + *acc += s; |
| 106 | + Some(start) |
| 107 | + }) |
| 108 | + .collect(); |
| 109 | + let sizes: Buffer<i32> = sizes.into_iter().collect(); |
| 110 | + ListViewArray::new(elements, offsets.into_array(), sizes.into_array(), validity).into_array() |
| 111 | +} |
| 112 | + |
| 113 | +/// A `FixedSizeList<i32, 8>` of `num_lists` lists, ~1/8 of them null. |
| 114 | +fn make_fsl(num_lists: usize) -> ArrayRef { |
| 115 | + let mut rng = StdRng::seed_from_u64(num_lists as u64); |
| 116 | + let total = (num_lists * BASE_LIST_SIZE) as i32; |
| 117 | + let elements = make_elements(total, false); |
| 118 | + let null_dist = Uniform::new(0u32, 8).unwrap(); |
| 119 | + let valid = (0..num_lists).map(|_| rng.sample(null_dist) != 0); |
| 120 | + let validity = Validity::Array(BoolArray::from_iter(valid).into_array()); |
| 121 | + FixedSizeListArray::new(elements, BASE_LIST_SIZE as u32, validity, num_lists).into_array() |
| 122 | +} |
| 123 | + |
| 124 | +/// Apply `list_sum(root())` and materialize the result. |
| 125 | +fn run(bencher: Bencher, array: ArrayRef) { |
| 126 | + let expr = list_sum(root()); |
| 127 | + bencher |
| 128 | + .with_inputs(|| (&array, SESSION.create_execution_ctx())) |
| 129 | + .bench_refs(|(array, ctx)| { |
| 130 | + array |
| 131 | + .clone() |
| 132 | + .apply(&expr) |
| 133 | + .unwrap() |
| 134 | + .execute::<Canonical>(ctx) |
| 135 | + .unwrap() |
| 136 | + }); |
| 137 | +} |
| 138 | + |
| 139 | +#[divan::bench] |
| 140 | +fn list_sum_small(bencher: Bencher) { |
| 141 | + run(bencher, make_list(SMALL, false)); |
| 142 | +} |
| 143 | + |
| 144 | +#[divan::bench] |
| 145 | +fn list_sum_medium(bencher: Bencher) { |
| 146 | + run(bencher, make_list(MEDIUM, false)); |
| 147 | +} |
| 148 | + |
| 149 | +#[divan::bench] |
| 150 | +fn list_sum_large(bencher: Bencher) { |
| 151 | + run(bencher, make_list(LARGE, false)); |
| 152 | +} |
| 153 | + |
| 154 | +#[divan::bench] |
| 155 | +fn list_sum_nullable_elements_medium(bencher: Bencher) { |
| 156 | + run(bencher, make_list(MEDIUM, true)); |
| 157 | +} |
| 158 | + |
| 159 | +#[divan::bench] |
| 160 | +fn list_sum_nullable_elements_large(bencher: Bencher) { |
| 161 | + run(bencher, make_list(LARGE, true)); |
| 162 | +} |
| 163 | + |
| 164 | +#[divan::bench] |
| 165 | +fn listview_sum_small(bencher: Bencher) { |
| 166 | + run(bencher, make_listview(SMALL)); |
| 167 | +} |
| 168 | + |
| 169 | +#[divan::bench] |
| 170 | +fn listview_sum_medium(bencher: Bencher) { |
| 171 | + run(bencher, make_listview(MEDIUM)); |
| 172 | +} |
| 173 | + |
| 174 | +#[divan::bench] |
| 175 | +fn listview_sum_large(bencher: Bencher) { |
| 176 | + run(bencher, make_listview(LARGE)); |
| 177 | +} |
| 178 | + |
| 179 | +#[divan::bench] |
| 180 | +fn fsl_sum_small(bencher: Bencher) { |
| 181 | + run(bencher, make_fsl(SMALL)); |
| 182 | +} |
| 183 | + |
| 184 | +#[divan::bench] |
| 185 | +fn fsl_sum_medium(bencher: Bencher) { |
| 186 | + run(bencher, make_fsl(MEDIUM)); |
| 187 | +} |
| 188 | + |
| 189 | +#[divan::bench] |
| 190 | +fn fsl_sum_large(bencher: Bencher) { |
| 191 | + run(bencher, make_fsl(LARGE)); |
| 192 | +} |
0 commit comments