|
| 1 | +use std::hint::black_box; |
| 2 | + |
| 3 | +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; |
| 4 | +use cubeorchestrator::query_message_parser::QueryResult; |
| 5 | +use cubeorchestrator::transport::JsRawColumnarData; |
| 6 | +use cubeshared::codegen::{ |
| 7 | + HttpColumnValue, HttpColumnValueArgs, HttpCommand, HttpMessage, HttpMessageArgs, HttpResultSet, |
| 8 | + HttpResultSetArgs, HttpRow, HttpRowArgs, |
| 9 | +}; |
| 10 | +use cubeshared::flatbuffers::FlatBufferBuilder; |
| 11 | + |
| 12 | +#[path = "common/mod.rs"] |
| 13 | +mod common; |
| 14 | +use common::{build_dataset, make_member_aliases, split_dim_measure, COLUMN_COUNTS, ROW_COUNTS}; |
| 15 | + |
| 16 | +/// Build a FlatBuffer `HttpMessage` payload mirroring CubeStore's wire format |
| 17 | +/// for `from_cubestore_fb` to parse. Cells are 16-character strings to give |
| 18 | +/// a realistic per-cell allocation cost. |
| 19 | +fn build_cubestore_fb_message(num_rows: usize, num_columns: usize) -> Vec<u8> { |
| 20 | + let mut builder = FlatBufferBuilder::new(); |
| 21 | + |
| 22 | + let column_names: Vec<_> = (0..num_columns) |
| 23 | + .map(|i| builder.create_string(&format!("column_{:02}", i))) |
| 24 | + .collect(); |
| 25 | + |
| 26 | + let mut rows_vec = Vec::with_capacity(num_rows); |
| 27 | + for row_idx in 0..num_rows { |
| 28 | + let mut values_vec = Vec::with_capacity(num_columns); |
| 29 | + for col_idx in 0..num_columns { |
| 30 | + let value_str = builder.create_string(&format!("r{:08}_c{:04}", row_idx, col_idx)); |
| 31 | + let col_value = HttpColumnValue::create( |
| 32 | + &mut builder, |
| 33 | + &HttpColumnValueArgs { |
| 34 | + string_value: Some(value_str), |
| 35 | + }, |
| 36 | + ); |
| 37 | + values_vec.push(col_value); |
| 38 | + } |
| 39 | + let values_vector = builder.create_vector(&values_vec); |
| 40 | + let row = HttpRow::create( |
| 41 | + &mut builder, |
| 42 | + &HttpRowArgs { |
| 43 | + values: Some(values_vector), |
| 44 | + }, |
| 45 | + ); |
| 46 | + rows_vec.push(row); |
| 47 | + } |
| 48 | + |
| 49 | + let columns_vector = builder.create_vector(&column_names); |
| 50 | + let rows_vector = builder.create_vector(&rows_vec); |
| 51 | + let result_set = HttpResultSet::create( |
| 52 | + &mut builder, |
| 53 | + &HttpResultSetArgs { |
| 54 | + columns: Some(columns_vector), |
| 55 | + rows: Some(rows_vector), |
| 56 | + }, |
| 57 | + ); |
| 58 | + |
| 59 | + let connection_id = builder.create_string("bench_connection"); |
| 60 | + let message = HttpMessage::create( |
| 61 | + &mut builder, |
| 62 | + &HttpMessageArgs { |
| 63 | + message_id: 1, |
| 64 | + command_type: HttpCommand::HttpResultSet, |
| 65 | + command: Some(result_set.as_union_value()), |
| 66 | + connection_id: Some(connection_id), |
| 67 | + }, |
| 68 | + ); |
| 69 | + |
| 70 | + builder.finish(message, None); |
| 71 | + builder.finished_data().to_vec() |
| 72 | +} |
| 73 | + |
| 74 | +fn bench_from_cubestore_fb(c: &mut Criterion) { |
| 75 | + let mut group = c.benchmark_group("QueryResult::from_cubestore_fb"); |
| 76 | + |
| 77 | + for &col_count in COLUMN_COUNTS { |
| 78 | + for &row_count in ROW_COUNTS { |
| 79 | + let msg = build_cubestore_fb_message(row_count, col_count); |
| 80 | + let id = format!("c{:02}_r{}", col_count, row_count); |
| 81 | + group.throughput(Throughput::Elements((row_count * col_count) as u64)); |
| 82 | + group.bench_with_input(BenchmarkId::from_parameter(id), &(), |b, _| { |
| 83 | + b.iter(|| { |
| 84 | + let result = |
| 85 | + QueryResult::from_cubestore_fb(black_box(&msg)).expect("from_cubestore_fb"); |
| 86 | + black_box(result); |
| 87 | + }); |
| 88 | + }); |
| 89 | + } |
| 90 | + } |
| 91 | + |
| 92 | + group.finish(); |
| 93 | +} |
| 94 | + |
| 95 | +fn bench_from_js_raw_data(c: &mut Criterion) { |
| 96 | + let mut group = c.benchmark_group("QueryResult::from_js_raw_data"); |
| 97 | + |
| 98 | + let combos: &[(usize, usize)] = &[(8, 10_000), (16, 10_000), (16, 100_000), (32, 100_000)]; |
| 99 | + |
| 100 | + for &(col_count, row_count) in combos { |
| 101 | + let (dim_count, measure_count) = split_dim_measure(col_count); |
| 102 | + let dimensions = make_member_aliases("dim", dim_count); |
| 103 | + let measures = make_member_aliases("measure", measure_count); |
| 104 | + |
| 105 | + let dataset = build_dataset(row_count, &dimensions, &measures, &[]); |
| 106 | + let payload = serde_json::to_vec(&dataset).expect("to_vec"); |
| 107 | + let payload_len = payload.len(); |
| 108 | + |
| 109 | + eprintln!( |
| 110 | + "from_js_raw_data: c{:02}_r{} payload_bytes={}", |
| 111 | + col_count, row_count, payload_len |
| 112 | + ); |
| 113 | + |
| 114 | + group.throughput(Throughput::Elements((row_count * col_count) as u64)); |
| 115 | + |
| 116 | + let id_param = format!("c{:02}_r{}", col_count, row_count); |
| 117 | + |
| 118 | + // Parse only: serde_json::from_slice into the wire type. |
| 119 | + group.bench_with_input(BenchmarkId::new("parse_only", &id_param), &(), |b, _| { |
| 120 | + b.iter(|| { |
| 121 | + let parsed: JsRawColumnarData = |
| 122 | + serde_json::from_slice(black_box(&payload)).expect("from_slice"); |
| 123 | + black_box(parsed); |
| 124 | + }); |
| 125 | + }); |
| 126 | + |
| 127 | + // End-to-end: parse + build into QueryResult — what the Neon bridge does. |
| 128 | + group.bench_with_input( |
| 129 | + BenchmarkId::new("parse_plus_build", &id_param), |
| 130 | + &(), |
| 131 | + |b, _| { |
| 132 | + b.iter(|| { |
| 133 | + let parsed: JsRawColumnarData = |
| 134 | + serde_json::from_slice(black_box(&payload)).expect("from_slice"); |
| 135 | + let built = QueryResult::from_js_raw_data(parsed).expect("from_js_raw_data"); |
| 136 | + black_box(built); |
| 137 | + }); |
| 138 | + }, |
| 139 | + ); |
| 140 | + } |
| 141 | + |
| 142 | + group.finish(); |
| 143 | +} |
| 144 | + |
| 145 | +criterion_group!(benches, bench_from_cubestore_fb, bench_from_js_raw_data); |
| 146 | +criterion_main!(benches); |
0 commit comments