diff --git a/benchmarks/src/dict.rs b/benchmarks/src/dict.rs index e04b5f816adc..067d111e5ca7 100644 --- a/benchmarks/src/dict.rs +++ b/benchmarks/src/dict.rs @@ -22,7 +22,9 @@ use arrow::array::{ArrayRef, DictionaryArray, Int32Array, ListArray, StringArray use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use arrow::record_batch::RecordBatch; +use clap::ValueEnum; use datafusion::common::DataFusionError; +use datafusion::datasource::MemTable; use datafusion::physical_plan::execute_stream; use datafusion::prelude::SessionContext; use datafusion_common::Result; @@ -44,6 +46,30 @@ pub struct RunOpt { #[clap(long)] pub query: Option, + /// How string group keys are represented at the aggregate input. + #[clap(long, value_enum, default_value_t = InputEncoding::Dictionary)] + pub input_encoding: InputEncoding, + + /// Override the generated number of distinct values. + #[clap(long)] + pub num_distinct: Option, + + /// Split the generated input into batches of this many rows. + #[clap(long)] + pub input_batch_size: Option, + + /// Pad generated strings to at least this many bytes. + #[clap(long)] + pub value_length: Option, + + /// Whether dictionary values are shared by all input batches. + #[clap(long, value_enum, default_value_t = DictionaryValueReuse::Shared)] + pub dictionary_value_reuse: DictionaryValueReuse, + + /// Cast dictionary group keys back to strings in the query output. + #[clap(long, default_value_t = false)] + pub normalize_output: bool, + /// Output path for the JSON benchmark summary. #[arg(short = 'o', long = "output")] pub output_path: Option, @@ -52,6 +78,37 @@ pub struct RunOpt { pub common: CommonOpt, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum InputEncoding { + /// Build dictionary arrays before query execution. + Dictionary, + /// Keep the group keys as plain string arrays. + Plain, + /// Keep the group keys as string view arrays. + View, + /// Build plain strings and dictionary-encode them in the query plan. + RuntimeDictionary, +} + +impl InputEncoding { + fn label(self) -> &'static str { + match self { + Self::Dictionary => "dictionary", + Self::Plain => "plain", + Self::View => "view", + Self::RuntimeDictionary => "runtime_dictionary", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum DictionaryValueReuse { + /// Slice one array so every batch shares the same dictionary values. + Shared, + /// Build an independent values array for every input batch. + PerBatch, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum DictValueType { Utf8, @@ -262,19 +319,37 @@ pub fn generate_dict( null_percent: NullPercent, size: usize, ) -> ArrayRef { - let num_distinct = cardinality.num_distinct(size); + generate_dict_with_offset(value_type, cardinality, None, None, null_percent, size, 0) +} + +fn generate_dict_with_offset( + value_type: DictValueType, + cardinality: Cardinality, + num_distinct: Option, + value_length: Option, + null_percent: NullPercent, + size: usize, + row_offset: usize, +) -> ArrayRef { + let num_distinct = num_distinct + .unwrap_or_else(|| cardinality.num_distinct(size)) + .clamp(1, i32::MAX as usize); let dict_values: ArrayRef = match value_type { DictValueType::Utf8 => { let strings: StringArray = (0..num_distinct) - .map(|i| format!("value_{i}")) + .map(|i| padded_value(format!("value_{i}"), value_length)) .collect::>() .into(); Arc::new(strings) } DictValueType::ListUtf8 => { let flat: StringArray = (0..num_distinct) - .flat_map(|i| (0..ITEMS_PER_VALUE).map(move |j| format!("value_{i}_{j}"))) + .flat_map(|i| { + (0..ITEMS_PER_VALUE).map(move |j| { + padded_value(format!("value_{i}_{j}"), value_length) + }) + }) .collect::>() .into(); @@ -296,10 +371,11 @@ pub fn generate_dict( let mut key_builder = Int32Array::builder(size); for i in 0..size { - if null_every.is_some_and(|n| i % n == 0) { + let row_index = row_offset + i; + if null_every.is_some_and(|n| row_index.is_multiple_of(n)) { key_builder.append_null(); } else { - key_builder.append_value((i % num_distinct) as i32); + key_builder.append_value((row_index % num_distinct) as i32); } } let keys = key_builder.finish(); @@ -309,6 +385,16 @@ pub fn generate_dict( .expect("valid dictionary array"), ) } + +fn padded_value(mut value: String, value_length: Option) -> String { + if let Some(value_length) = value_length { + value.extend(std::iter::repeat_n( + 'x', + value_length.saturating_sub(value.len()), + )); + } + value +} impl RunOpt { pub async fn run(self) -> Result<()> { println!( @@ -337,7 +423,30 @@ impl RunOpt { for query_id in query_range { let query = &DICTIONARY_QUERIES[query_id - 1]; - benchmark_run.start_new_case(query.name); + let case_name = if self.input_encoding == InputEncoding::Dictionary + && !self.normalize_output + && self.dictionary_value_reuse == DictionaryValueReuse::Shared + { + query.name.to_string() + } else { + let output = if self.normalize_output { + "_normalized" + } else { + "" + }; + let dictionary_values = + if self.dictionary_value_reuse == DictionaryValueReuse::PerBatch { + "_per_batch_values" + } else { + "" + }; + format!( + "{}_{}{output}{dictionary_values}", + query.name, + self.input_encoding.label() + ) + }; + benchmark_run.start_new_case(&case_name); let query_run = self.benchmark_query(query, &ctx).await; match query_run { @@ -364,15 +473,21 @@ impl RunOpt { query: &DictionaryQuery, ctx: &SessionContext, ) -> Result> { - let batch = self.make_record_batch(query)?; + let batches = self.make_record_batches(query)?; + let schema = batches.first().expect("at least one input batch").schema(); ctx.deregister_table("test_data")?; - ctx.register_batch("test_data", batch)?; + ctx.register_table( + "test_data", + Arc::new(MemTable::try_new(schema, vec![batches])?), + )?; + + let sql = self.query_sql(query); let mut query_results = vec![]; for i in 0..self.common.iterations { let start = Instant::now(); let row_count = - Self::execute_query_without_result_buffering(query.sql, ctx).await?; + Self::execute_query_without_result_buffering(&sql, ctx).await?; let elapsed = start.elapsed(); println!( "Query '{}' iteration {i} returned {row_count} rows in {elapsed:?}", @@ -384,16 +499,137 @@ impl RunOpt { Ok(query_results) } - fn make_record_batch(&self, query: &DictionaryQuery) -> Result { + fn query_sql(&self, query: &DictionaryQuery) -> String { + if self.input_encoding == InputEncoding::Plain + || (matches!( + self.input_encoding, + InputEncoding::Dictionary | InputEncoding::View + ) && !self.normalize_output) + { + return query.sql.to_string(); + } + + let group_col = match self.input_encoding { + InputEncoding::Dictionary | InputEncoding::View => "dict_col", + InputEncoding::RuntimeDictionary => { + "arrow_cast(dict_col, 'Dictionary(Int32, Utf8)')" + } + InputEncoding::Plain => unreachable!(), + }; + let select_col = if self.normalize_output { + format!("arrow_cast({group_col}, 'Utf8') AS dict_col") + } else { + group_col.to_string() + }; + + if query.col2.is_some() { + let group_col2 = match self.input_encoding { + InputEncoding::Dictionary | InputEncoding::View => "dict_col2", + InputEncoding::RuntimeDictionary => { + "arrow_cast(dict_col2, 'Dictionary(Int32, Utf8)')" + } + InputEncoding::Plain => unreachable!(), + }; + let select_col2 = if self.normalize_output { + format!("arrow_cast({group_col2}, 'Utf8') AS dict_col2") + } else { + group_col2.to_string() + }; + format!( + "SELECT {select_col}, {select_col2}, COUNT(*) FROM test_data GROUP BY {group_col}, {group_col2}" + ) + } else { + format!("SELECT {select_col}, COUNT(*) FROM test_data GROUP BY {group_col}") + } + } + + fn make_record_batches(&self, query: &DictionaryQuery) -> Result> { let size = self.num_rows; + if size == 0 { + return exec_err!("num_rows must be greater than zero"); + } + if let Some(num_distinct) = self.num_distinct { + if num_distinct == 0 { + return exec_err!("num_distinct must be greater than zero"); + } + if num_distinct > size { + return exec_err!( + "num_distinct ({num_distinct}) must not exceed num_rows ({size})" + ); + } + if num_distinct > i32::MAX as usize { + return exec_err!( + "num_distinct must fit in the Int32 dictionary key type" + ); + } + } + let input_batch_size = self.input_batch_size.unwrap_or(size); + if input_batch_size == 0 { + return exec_err!("input_batch_size must be greater than zero"); + } + if self.input_encoding != InputEncoding::Dictionary + && self.dictionary_value_reuse == DictionaryValueReuse::PerBatch + { + return exec_err!( + "dictionary_value_reuse=per-batch requires input_encoding=dictionary" + ); + } + + if self.input_encoding == InputEncoding::Dictionary + && self.dictionary_value_reuse == DictionaryValueReuse::PerBatch + { + return (0..size) + .step_by(input_batch_size) + .map(|offset| { + self.make_record_batch( + query, + input_batch_size.min(size - offset), + offset, + ) + }) + .collect(); + } - let col1 = generate_dict( + let batch = self.make_record_batch(query, size, 0)?; + Ok((0..size) + .step_by(input_batch_size) + .map(|offset| batch.slice(offset, input_batch_size.min(size - offset))) + .collect()) + } + + fn make_record_batch( + &self, + query: &DictionaryQuery, + size: usize, + row_offset: usize, + ) -> Result { + let num_distinct = |params: &DictColParams| { + Some( + self.num_distinct + .unwrap_or_else(|| params.cardinality.num_distinct(self.num_rows)), + ) + }; + + let col1 = generate_dict_with_offset( query.col.value_type, query.col.cardinality, + num_distinct(&query.col), + self.value_length, query.col.null_percent, size, + row_offset, ); + let col1 = match self.input_encoding { + InputEncoding::Dictionary => col1, + InputEncoding::Plain | InputEncoding::RuntimeDictionary => { + arrow::compute::cast(col1.as_ref(), &DataType::Utf8)? + } + InputEncoding::View => { + arrow::compute::cast(col1.as_ref(), &DataType::Utf8View)? + } + }; + let (schema, columns): (Schema, Vec) = match &query.col2 { None => { let schema = Schema::new(vec![Field::new( @@ -404,12 +640,24 @@ impl RunOpt { (schema, vec![col1]) } Some(col2_params) => { - let col2 = generate_dict( + let col2 = generate_dict_with_offset( col2_params.value_type, col2_params.cardinality, + num_distinct(col2_params), + self.value_length, col2_params.null_percent, size, + row_offset, ); + let col2 = match self.input_encoding { + InputEncoding::Dictionary => col2, + InputEncoding::Plain | InputEncoding::RuntimeDictionary => { + arrow::compute::cast(col2.as_ref(), &DataType::Utf8)? + } + InputEncoding::View => { + arrow::compute::cast(col2.as_ref(), &DataType::Utf8View)? + } + }; let schema = Schema::new(vec![ Field::new("dict_col", col1.data_type().clone(), true), Field::new("dict_col2", col2.data_type().clone(), true), @@ -435,3 +683,216 @@ impl RunOpt { Ok(row_count) } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Int64Array}; + use std::collections::BTreeMap; + + fn run_opt(input_encoding: InputEncoding) -> RunOpt { + RunOpt { + num_rows: 10, + query: Some(1), + input_encoding, + num_distinct: Some(3), + input_batch_size: Some(4), + value_length: Some(32), + dictionary_value_reuse: DictionaryValueReuse::Shared, + normalize_output: true, + output_path: None, + common: CommonOpt { + iterations: 1, + partitions: Some(1), + batch_size: None, + mem_pool_type: "fair".to_string(), + memory_limit: None, + sort_spill_reservation_bytes: None, + debug: false, + simulate_latency: false, + }, + } + } + + #[test] + fn dictionary_batches_share_values_and_preserve_padding() { + let batches = run_opt(InputEncoding::Dictionary) + .make_record_batches(&DICTIONARY_QUERIES[0]) + .unwrap(); + + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![4, 4, 2] + ); + + let first = batches[0].column(0); + let second = batches[1].column(0); + let first = first + .as_any() + .downcast_ref::>() + .unwrap(); + let second = second + .as_any() + .downcast_ref::>() + .unwrap(); + assert!(Arc::ptr_eq(first.values(), second.values())); + + let values = first + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.len(), 3); + assert!(values.iter().flatten().all(|value| value.len() == 32)); + } + + #[test] + fn query_sql_keeps_user_visible_string_type() { + let query = &DICTIONARY_QUERIES[0]; + assert_eq!(run_opt(InputEncoding::Plain).query_sql(query), query.sql); + + let mut default_dictionary = run_opt(InputEncoding::Dictionary); + default_dictionary.normalize_output = false; + assert_eq!(default_dictionary.query_sql(query), query.sql); + + assert_eq!( + run_opt(InputEncoding::Dictionary).query_sql(query), + "SELECT arrow_cast(dict_col, 'Utf8') AS dict_col, COUNT(*) FROM test_data GROUP BY dict_col" + ); + assert_eq!( + run_opt(InputEncoding::View).query_sql(query), + "SELECT arrow_cast(dict_col, 'Utf8') AS dict_col, COUNT(*) FROM test_data GROUP BY dict_col" + ); + assert_eq!( + run_opt(InputEncoding::RuntimeDictionary).query_sql(query), + "SELECT arrow_cast(arrow_cast(dict_col, 'Dictionary(Int32, Utf8)'), 'Utf8') AS dict_col, COUNT(*) FROM test_data GROUP BY arrow_cast(dict_col, 'Dictionary(Int32, Utf8)')" + ); + } + + #[test] + fn per_batch_dictionaries_keep_global_row_pattern() { + let mut opt = run_opt(InputEncoding::Dictionary); + opt.dictionary_value_reuse = DictionaryValueReuse::PerBatch; + let batches = opt.make_record_batches(&DICTIONARY_QUERIES[0]).unwrap(); + + let first = batches[0] + .column(0) + .as_any() + .downcast_ref::>() + .unwrap(); + let second = batches[1] + .column(0) + .as_any() + .downcast_ref::>() + .unwrap(); + assert!(!Arc::ptr_eq(first.values(), second.values())); + assert_eq!(first.keys().value(0), 0); + assert_eq!(second.keys().value(0), 1); + } + + #[test] + fn rejects_zero_input_batch_size() { + let mut opt = run_opt(InputEncoding::Dictionary); + opt.input_batch_size = Some(0); + let err = opt.make_record_batches(&DICTIONARY_QUERIES[0]).unwrap_err(); + assert!( + err.to_string() + .contains("input_batch_size must be greater than zero") + ); + } + + #[test] + fn rejects_invalid_distinct_count_and_dictionary_reuse() { + let query = &DICTIONARY_QUERIES[0]; + + let mut zero_distinct = run_opt(InputEncoding::Dictionary); + zero_distinct.num_distinct = Some(0); + assert!( + zero_distinct + .make_record_batches(query) + .unwrap_err() + .to_string() + .contains("num_distinct must be greater than zero") + ); + + let mut excess_distinct = run_opt(InputEncoding::Dictionary); + excess_distinct.num_distinct = Some(excess_distinct.num_rows + 1); + assert!( + excess_distinct + .make_record_batches(query) + .unwrap_err() + .to_string() + .contains("must not exceed num_rows") + ); + + let mut plain_per_batch = run_opt(InputEncoding::Plain); + plain_per_batch.dictionary_value_reuse = DictionaryValueReuse::PerBatch; + assert!( + plain_per_batch + .make_record_batches(query) + .unwrap_err() + .to_string() + .contains("requires input_encoding=dictionary") + ); + } + + async fn grouped_counts(opt: &RunOpt) -> BTreeMap { + let query = &DICTIONARY_QUERIES[0]; + let batches = opt.make_record_batches(query).unwrap(); + let schema = batches[0].schema(); + let ctx = SessionContext::new(); + ctx.register_table( + "test_data", + Arc::new(MemTable::try_new(schema, vec![batches]).unwrap()), + ) + .unwrap(); + + let batches = ctx + .sql(&opt.query_sql(query)) + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut counts = BTreeMap::new(); + for batch in batches { + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + counts.insert(keys.value(row).to_string(), values.value(row)); + } + } + counts + } + + #[tokio::test] + async fn encodings_produce_identical_grouped_counts() { + let plain = run_opt(InputEncoding::Plain); + let shared = run_opt(InputEncoding::Dictionary); + let mut per_batch = run_opt(InputEncoding::Dictionary); + per_batch.dictionary_value_reuse = DictionaryValueReuse::PerBatch; + let runtime = run_opt(InputEncoding::RuntimeDictionary); + let view = run_opt(InputEncoding::View); + + let expected = grouped_counts(&plain).await; + assert_eq!( + expected.values().copied().collect::>(), + vec![4, 3, 3] + ); + assert_eq!(grouped_counts(&shared).await, expected); + assert_eq!(grouped_counts(&per_batch).await, expected); + assert_eq!(grouped_counts(&runtime).await, expected); + assert_eq!(grouped_counts(&view).await, expected); + } +}