diff --git a/datafusion/core/benches/data_utils/mod.rs b/datafusion/core/benches/data_utils/mod.rs index 728c6490c72bd..213fb72760608 100644 --- a/datafusion/core/benches/data_utils/mod.rs +++ b/datafusion/core/benches/data_utils/mod.rs @@ -26,13 +26,15 @@ use arrow::datatypes::Int32Type; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::MemTable; use datafusion::error::Result; -use datafusion_common::DataFusionError; +use datafusion_common::{ + DataFusionError, + utils::hex::{HexCase, encode_bytes_to_slice}, +}; use rand::prelude::IndexedRandom; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use rand_distr::Distribution; use rand_distr::{Normal, Pareto}; -use std::fmt::Write; use std::sync::Arc; /// create an in-memory table given the partition len, array len, and batch size, @@ -217,21 +219,15 @@ pub(crate) fn make_data( let mut cur_time = 16909000000000i64; for _ in 0..partition_cnt { // Choose the appropriate builder based on use_view. + let sample_cnt = sample_cnt as usize; let mut id_builder = if use_view { - TraceIdBuilder::Utf8View(StringViewBuilder::new()) + TraceIdBuilder::Utf8View(StringViewBuilder::with_capacity(sample_cnt)) } else { - TraceIdBuilder::Utf8(StringBuilder::new()) + TraceIdBuilder::Utf8(StringBuilder::with_capacity(sample_cnt, 32)) }; - let mut ts_builder = Int64Builder::new(); - let gen_id = |rng: &mut rand::rngs::SmallRng| { - rng.random::<[u8; 16]>() - .iter() - .fold(String::new(), |mut output, b| { - let _ = write!(output, "{b:02X}"); - output - }) - }; + let mut ts_builder = Int64Builder::with_capacity(sample_cnt); + let gen_id = |rng: &mut rand::rngs::SmallRng| rng.random::<[u8; 16]>(); let gen_sample_cnt = |mut rng: &mut rand::rngs::SmallRng| pareto.sample(&mut rng).ceil() as u32; let mut group_ids = (0..simultaneous_group_cnt) @@ -240,6 +236,7 @@ pub(crate) fn make_data( let mut group_sample_cnts = (0..simultaneous_group_cnt) .map(|_| gen_sample_cnt(&mut rng)) .collect::>(); + for _ in 0..sample_cnt { let random_index = rng.random_range(0..simultaneous_group_cnt); let trace_id = &mut group_ids[random_index]; @@ -250,7 +247,12 @@ pub(crate) fn make_data( *sample_cnt = gen_sample_cnt(&mut rng); } - id_builder.append_value(trace_id); + let mut id_out = [0; 32]; + encode_bytes_to_slice(trace_id, HexCase::Upper, &mut id_out).unwrap(); + id_builder.append_value({ + // SAFETY: `id_out` holds only ASCII hex digits, which are valid UTF-8. + unsafe { str::from_utf8_unchecked(&id_out) } + }); ts_builder.append_value(cur_time); if asc { diff --git a/datafusion/core/benches/topk_aggregate.rs b/datafusion/core/benches/topk_aggregate.rs index 8fc6c954caa9c..6873a96260c96 100644 --- a/datafusion/core/benches/topk_aggregate.rs +++ b/datafusion/core/benches/topk_aggregate.rs @@ -35,6 +35,12 @@ use tokio::runtime::Runtime; const LIMIT: usize = 10; +const LARGE_LIMIT: usize = 50000; + +const PARTITIONS: i32 = 10; + +const SAMPLES: i32 = 1_000_000; + /// Create deterministic data for DISTINCT benchmarks with predictable trace_ids /// This ensures consistent results across benchmark runs fn make_distinct_data( @@ -80,7 +86,7 @@ fn create_context( asc: bool, use_topk: bool, use_view: bool, -) -> Result { +) -> SessionContext { let (schema, parts) = make_data(partition_cnt, sample_cnt, asc, use_view).unwrap(); let mem_table = Arc::new(MemTable::try_new(schema, parts).unwrap()); @@ -89,9 +95,9 @@ fn create_context( let opts = cfg.options_mut(); opts.optimizer.enable_topk_aggregation = use_topk; let ctx = SessionContext::new_with_config(cfg); - let _ = ctx.register_table("traces", mem_table)?; + ctx.register_table("traces", mem_table).unwrap(); - Ok(ctx) + ctx } fn create_context_distinct( @@ -113,32 +119,41 @@ fn create_context_distinct( Ok(ctx) } -fn run(rt: &Runtime, ctx: SessionContext, limit: usize, use_topk: bool, asc: bool) { - black_box(rt.block_on(async { aggregate(ctx, limit, use_topk, asc).await })).unwrap(); +fn run( + rt: &Runtime, + ctx: &SessionContext, + limit: usize, + use_topk: bool, +) -> Vec { + black_box(rt.block_on(async { aggregate(ctx, limit, use_topk).await })).unwrap() } -fn run_string(rt: &Runtime, ctx: SessionContext, limit: usize, use_topk: bool) { +fn run_string( + rt: &Runtime, + ctx: &SessionContext, + limit: usize, + use_topk: bool, +) -> Vec { black_box(rt.block_on(async { aggregate_string(ctx, limit, use_topk).await })) - .unwrap(); + .unwrap() } fn run_distinct( rt: &Runtime, - ctx: SessionContext, + ctx: &SessionContext, limit: usize, use_topk: bool, asc: bool, -) { +) -> Vec { black_box(rt.block_on(async { aggregate_distinct(ctx, limit, use_topk, asc).await })) - .unwrap(); + .unwrap() } async fn aggregate( - ctx: SessionContext, + ctx: &SessionContext, limit: usize, use_topk: bool, - asc: bool, -) -> Result<()> { +) -> Result> { let sql = format!( "select max(timestamp_ms) from traces group by trace_id order by max(timestamp_ms) desc limit {limit};" ); @@ -151,11 +166,17 @@ async fn aggregate( ); let batches = collect(plan, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), limit); + + Ok(batches) +} + +fn validate_aggregate(batches: &[RecordBatch], asc: bool) -> Result<()> { assert_eq!(batches.len(), 1); let batch = batches.first().unwrap(); assert_eq!(batch.num_rows(), LIMIT); - let actual = format!("{}", pretty_format_batches(&batches)?).to_lowercase(); + let actual = pretty_format_batches(batches)?.to_string().to_lowercase(); let expected_asc = r#" +--------------------------+ | max(traces.timestamp_ms) | @@ -184,7 +205,7 @@ async fn aggregate( /// This tests grouping by a numeric column (timestamp_ms) and aggregating /// a string column (trace_id) with Utf8 or Utf8View data types. async fn aggregate_string( - ctx: SessionContext, + ctx: &SessionContext, limit: usize, use_topk: bool, ) -> Result> { @@ -200,19 +221,17 @@ async fn aggregate_string( ); let batches = collect(plan, ctx.task_ctx()).await?; - assert_eq!(batches.len(), 1); - let batch = batches.first().unwrap(); - assert_eq!(batch.num_rows(), LIMIT); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), limit); Ok(batches) } async fn aggregate_distinct( - ctx: SessionContext, + ctx: &SessionContext, limit: usize, use_topk: bool, asc: bool, -) -> Result<()> { +) -> Result> { let order_direction = if asc { "asc" } else { "desc" }; let sql = format!( "select id from traces group by id order by id {order_direction} limit {limit};" @@ -225,11 +244,17 @@ async fn aggregate_distinct( use_topk ); let batches = collect(plan, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), limit); + + Ok(batches) +} + +fn validate_aggregate_distinct(batches: &[RecordBatch], asc: bool) -> Result<()> { assert_eq!(batches.len(), 1); let batch = batches.first().unwrap(); assert_eq!(batch.num_rows(), LIMIT); - let actual = format!("{}", pretty_format_batches(&batches)?).to_lowercase(); + let actual = pretty_format_batches(batches)?.to_string().to_lowercase(); let expected_asc = r#" +----+ @@ -300,14 +325,12 @@ struct StringCase { fn assert_utf8_utf8view_match( rt: &Runtime, - partitions: i32, - samples: i32, + ctx_utf8: &SessionContext, + ctx_view: &SessionContext, limit: usize, asc: bool, use_topk: bool, ) { - let ctx_utf8 = create_context(partitions, samples, asc, use_topk, false).unwrap(); - let ctx_view = create_context(partitions, samples, asc, use_topk, true).unwrap(); let batches_utf8 = rt .block_on(aggregate_string(ctx_utf8, limit, use_topk)) .unwrap(); @@ -326,11 +349,17 @@ fn assert_string_results_match( rt: &Runtime, partitions: i32, samples: i32, - limit: usize, + limits: &[usize], ) { for asc in [false, true] { for use_topk in [false, true] { - assert_utf8_utf8view_match(rt, partitions, samples, limit, asc, use_topk); + let ctx_utf8 = create_context(partitions, samples, asc, use_topk, false); + let ctx_view = create_context(partitions, samples, asc, use_topk, true); + for &limit in limits { + assert_utf8_utf8view_match( + rt, &ctx_utf8, &ctx_view, limit, asc, use_topk, + ); + } } } } @@ -341,21 +370,20 @@ fn assert_string_results_match( )] fn criterion_benchmark(c: &mut Criterion) { let rt = Runtime::new().unwrap(); - let limit = LIMIT; - let partitions = 10; - let samples = 1_000_000; + let partitions = PARTITIONS; + let samples = SAMPLES; let total_rows = partitions * samples; // Numeric aggregate benchmarks let numeric_cases = &[ BenchCase { - name_tpl: "aggregate {rows} time-series rows", + name_tpl: "k={limit} aggregate {rows} time-series rows", asc: false, use_topk: false, use_view: false, }, BenchCase { - name_tpl: "aggregate {rows} worst-case rows", + name_tpl: "k={limit} aggregate {rows} worst-case rows", asc: true, use_topk: false, use_view: false, @@ -386,19 +414,25 @@ fn criterion_benchmark(c: &mut Criterion) { }, ]; for case in numeric_cases { - let name = case - .name_tpl - .replace("{rows}", &total_rows.to_string()) - .replace("{limit}", &limit.to_string()); let ctx = - create_context(partitions, samples, case.asc, case.use_topk, case.use_view) - .unwrap(); + create_context(partitions, samples, case.asc, case.use_topk, case.use_view); + + let name = case.name_tpl.replace("{rows}", &total_rows.to_string()); c.bench_function(&name, |b| { - b.iter(|| run(&rt, ctx.clone(), limit, case.use_topk, case.asc)) + b.iter(|| { + let batches = run(&rt, &ctx, LIMIT, case.use_topk); + validate_aggregate(&batches, case.asc).unwrap(); + batches + }) + }); + + let name = case.name_tpl.replace("{rows}", &total_rows.to_string()); + c.bench_function(&name, |b| { + b.iter(|| run(&rt, &ctx, LARGE_LIMIT, case.use_topk)) }); } - assert_string_results_match(&rt, partitions, samples, limit); + assert_string_results_match(&rt, partitions, samples, &[LIMIT, LARGE_LIMIT]); let string_cases = &[ StringCase { @@ -449,19 +483,23 @@ fn criterion_benchmark(c: &mut Criterion) { "time-series" }; let type_label = if case.use_view { "Utf8View" } else { "Utf8" }; - let name = if case.use_topk { - format!( - "top k={limit} string aggregate {total_rows} {scenario} rows [{type_label}]" - ) - } else { - format!("string aggregate {total_rows} {scenario} rows [{type_label}]") - }; let ctx = - create_context(partitions, samples, case.asc, case.use_topk, case.use_view) - .unwrap(); - c.bench_function(&name, |b| { - b.iter(|| run_string(&rt, ctx.clone(), limit, case.use_topk)) - }); + create_context(partitions, samples, case.asc, case.use_topk, case.use_view); + + for limit in [LIMIT, LARGE_LIMIT] { + let name = if case.use_topk { + format!( + "top k={limit} string aggregate {total_rows} {scenario} rows [{type_label}]" + ) + } else { + format!( + "k={limit} string aggregate {total_rows} {scenario} rows [{type_label}]" + ) + }; + c.bench_function(&name, |b| { + b.iter(|| run_string(&rt, &ctx, limit, case.use_topk)) + }); + } } // DISTINCT benchmarks @@ -470,9 +508,21 @@ fn criterion_benchmark(c: &mut Criterion) { let topk_label = if use_topk { "TopK" } else { "no TopK" }; for asc in [false, true] { let dir = if asc { "asc" } else { "desc" }; - let name = format!("distinct {total_rows} rows {dir} [{topk_label}]"); + let name = + format!("top k={LIMIT} distinct {total_rows} rows {dir} [{topk_label}]"); + c.bench_function(&name, |b| { + b.iter(|| { + let batches = run_distinct(&rt, &ctx, LIMIT, use_topk, asc); + validate_aggregate_distinct(&batches, asc).unwrap(); + batches + }) + }); + + let name = format!( + "top k={LARGE_LIMIT} distinct {total_rows} rows {dir} [{topk_label}]" + ); c.bench_function(&name, |b| { - b.iter(|| run_distinct(&rt, ctx.clone(), limit, use_topk, asc)) + b.iter(|| run_distinct(&rt, &ctx, LARGE_LIMIT, use_topk, asc)) }); } } diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 030614eed02e5..6938b61145fb2 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -499,7 +499,7 @@ impl TopK { Some(max_row) if row.as_ref() >= max_row.row() => {} // don't yet have k items or new item is lower than the currently k low values None | Some(_) => { - self.heap.add(batch_entry, row, index); + self.heap.add(batch_entry, row.as_ref(), index, false); replacements += 1; } } @@ -857,6 +857,31 @@ impl TopKMetrics { } } +enum TopKData { + Sorted(Vec), + Raw(BinaryHeap), +} + +impl TopKData { + fn into_sorted(self) -> Vec { + match self { + TopKData::Sorted(sorted) => sorted, + TopKData::Raw(raw) => { + // TODO: sort_unstable is a lot faster, but changes order slightly... + // when TopKRow Ord takes batch_id and index into account, + // we get the same result for both + //let mut vec = raw.into_vec(); + //vec.sort_unstable(); + //vec + + // BinaryHeap::into_sorted_vec is slow, + // see https://github.com/rust-lang/rust/issues/115357 + raw.into_sorted_vec() + } + } + } +} + /// This structure keeps at most the *smallest* k items, using the /// [arrow::row] format for sort keys. While it is called "topK" for /// values like `1, 2, 3, 4, 5` the "top 3" really means the @@ -920,14 +945,14 @@ impl TopKHeap { fn add( &mut self, batch_entry: &mut RecordBatchEntry, - row: impl AsRef<[u8]>, + row: &[u8], index: usize, + store_eviction: bool, ) -> Option { let batch_id = batch_entry.id; batch_entry.uses += 1; assert!(self.inner.len() <= self.k); - let row = row.as_ref(); // Reuse storage for evicted item if possible if self.inner.len() == self.k { @@ -939,18 +964,23 @@ impl TopKHeap { // cross-batch evictions, or directly from `batch_entry` when // a row evicts another row from the same in-flight batch // (entry not yet registered in the store). - let evicted_batch = if prev_min.batch_id == batch_entry.id { - batch_entry.batch.clone() + let evicted = if store_eviction { + let evicted_batch = if prev_min.batch_id == batch_entry.id { + batch_entry.batch.clone() + } else { + self.store + .get(prev_min.batch_id) + .expect("evicted row's batch must be present in the store") + .batch + .clone() + }; + Some(EvictedRow { + batch: evicted_batch, + index: prev_min.index, + row_bytes: prev_min.row.clone(), + }) } else { - self.store - .get(prev_min.batch_id) - .map(|entry| entry.batch.clone()) - .expect("evicted row's batch must be present in the store") - }; - let evicted = EvictedRow { - batch: evicted_batch, - index: prev_min.index, - row_bytes: prev_min.row.clone(), + None }; // Update batch use @@ -967,7 +997,7 @@ impl TopKHeap { self.owned_bytes += prev_min.owned_size(); - Some(evicted) + evicted } else { let new_row = TopKRow::new(row, batch_id, index); self.owned_bytes += new_row.owned_size(); @@ -986,17 +1016,17 @@ impl TopKHeap { /// Returns the values stored in this heap, from values low to /// high, as a single [`RecordBatch`], and a sorted vec of the /// current heap's contents - fn emit_with_state(&mut self) -> Result<(Option, Vec)> { - // generate sorted rows - let topk_rows = std::mem::take(&mut self.inner).into_sorted_vec(); - + fn emit_with_state(&mut self) -> Result<(Option, TopKData)> { + let topk_rows = TopKData::Raw(std::mem::take(&mut self.inner)); if self.store.is_empty() { return Ok((None, topk_rows)); } + // sort rows after check + let topk_rows = topk_rows.into_sorted(); // Collect the batches into a vec and store the "batch_id -> array_pos" mapping, to then // build the `indices` vec below. This is needed since the batch ids are not continuous. - let mut record_batches = Vec::new(); + let mut record_batches = Vec::with_capacity(self.store.batches.len()); let mut batch_id_array_pos = HashMap::new(); for (array_pos, (batch_id, batch)) in self.store.batches.iter().enumerate() { record_batches.push(&batch.batch); @@ -1014,7 +1044,7 @@ impl TopKHeap { // them together into a single new batch let new_batch = interleave_record_batch(&record_batches, &indices)?; - Ok((Some(new_batch), topk_rows)) + Ok((Some(new_batch), TopKData::Sorted(topk_rows))) } /// Compact this heap, rewriting all stored batches into a single @@ -1040,10 +1070,11 @@ impl TopKHeap { // batches that have a high usage ratio already // Note: new batch is in the same order as inner - let (new_batch, mut topk_rows) = self.emit_with_state()?; + let (new_batch, topk_rows) = self.emit_with_state()?; let Some(new_batch) = new_batch else { return Ok(()); }; + let mut topk_rows = topk_rows.into_sorted(); // clear all old entries in store (this invalidates all // store_ids in `inner`) @@ -1080,7 +1111,7 @@ impl TopKHeap { /// also be primitive values) /// /// Reuses allocations to minimize runtime overhead of creating new Vecs -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] struct TopKRow { /// the value of the sort key for this row. This contains the /// bytes that could be stored in `OwnedRow` but uses `Vec` to @@ -1386,7 +1417,7 @@ impl PartitionedTopK { match heap.max() { Some(max_row) if row.as_ref() >= max_row.row() => {} None | Some(_) => { - heap.add(&mut entry, row, sub_idx); + heap.add(&mut entry, row.as_ref(), sub_idx, false); replacements += 1; } } @@ -1685,8 +1716,12 @@ impl PartitionedTopKRank { batch: evicted_batch, index: evicted_index, row_bytes: evicted_bytes, - }) = state.heap.add(entry_ref, row, orig_idx as usize) - { + }) = state.heap.add( + entry_ref, + row.as_ref(), + orig_idx as usize, + true, + ) { // Compare the new boundary (post-eviction heap // top) against the evicted row's bytes — both // already in encoded form, no clones needed.