This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 7df3e5cd11 Add benchmark for memory-limited aggregation (#13090)
7df3e5cd11 is described below
commit 7df3e5cd11f63226b90783564ae7268ee2512ec1
Author: Yongting You <[email protected]>
AuthorDate: Sat Oct 26 18:59:09 2024 +0800
Add benchmark for memory-limited aggregation (#13090)
* Adding benchmark for external aggregation
* comments
---
benchmarks/README.md | 28 ++
benchmarks/bench.sh | 25 +-
benchmarks/src/bin/external_aggr.rs | 390 ++++++++++++++++++++++++++++
datafusion/execution/src/memory_pool/mod.rs | 12 +-
4 files changed, 450 insertions(+), 5 deletions(-)
diff --git a/benchmarks/README.md b/benchmarks/README.md
index a12662ccb8..a9aa1afb97 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -352,6 +352,34 @@ This benchmarks is derived from the [TPC-H][1] version
[2]: https://github.com/databricks/tpch-dbgen.git,
[2.17.1]:
https://www.tpc.org/tpc_documents_current_versions/pdf/tpc-h_v2.17.1.pdf
+## External Aggregation
+
+Run the benchmark for aggregations with limited memory.
+
+When the memory limit is exceeded, the aggregation intermediate results will
be spilled to disk, and finally read back with sort-merge.
+
+External aggregation benchmarks run several aggregation queries with different
memory limits, on TPCH `lineitem` table. Queries can be found in
[`external_aggr.rs`](src/bin/external_aggr.rs).
+
+This benchmark is inspired by [DuckDB's external aggregation
paper](https://hannes.muehleisen.org/publications/icde2024-out-of-core-kuiper-boncz-muehleisen.pdf),
specifically Section VI.
+
+### External Aggregation Example Runs
+1. Run all queries with predefined memory limits:
+```bash
+# Under 'benchmarks/' directory
+cargo run --release --bin external_aggr -- benchmark -n 4 --iterations 3 -p
'....../data/tpch_sf1' -o '/tmp/aggr.json'
+```
+
+2. Run a query with specific memory limit:
+```bash
+cargo run --release --bin external_aggr -- benchmark -n 4 --iterations 3 -p
'....../data/tpch_sf1' -o '/tmp/aggr.json' --query 1 --memory-limit 30M
+```
+
+3. Run all queries with `bench.sh` script:
+```bash
+./bench.sh data external_aggr
+./bench.sh run external_aggr
+```
+
# Older Benchmarks
diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh
index fc10cc5afc..47c5d12616 100755
--- a/benchmarks/bench.sh
+++ b/benchmarks/bench.sh
@@ -78,6 +78,7 @@ sort: Benchmark of sorting speed
clickbench_1: ClickBench queries against a single parquet file
clickbench_partitioned: ClickBench queries against a partitioned (100 files)
parquet
clickbench_extended: ClickBench \"inspired\" queries against a single
parquet (DataFusion specific)
+external_aggr: External aggregation benchmark
**********
* Supported Configuration (Environment Variables)
@@ -170,6 +171,10 @@ main() {
imdb)
data_imdb
;;
+ external_aggr)
+ # same data as for tpch
+ data_tpch "1"
+ ;;
*)
echo "Error: unknown benchmark '$BENCHMARK' for data
generation"
usage
@@ -212,6 +217,7 @@ main() {
run_clickbench_partitioned
run_clickbench_extended
run_imdb
+ run_external_aggr
;;
tpch)
run_tpch "1"
@@ -243,6 +249,9 @@ main() {
imdb)
run_imdb
;;
+ external_aggr)
+ run_external_aggr
+ ;;
*)
echo "Error: unknown benchmark '$BENCHMARK' for run"
usage
@@ -524,7 +533,21 @@ run_imdb() {
$CARGO_COMMAND --bin imdb -- benchmark datafusion --iterations 5 --path
"${IMDB_DIR}" --prefer_hash_join "${PREFER_HASH_JOIN}" --format parquet -o
"${RESULTS_FILE}"
}
-
+# Runs the external aggregation benchmark
+run_external_aggr() {
+ # Use TPC-H SF1 dataset
+ TPCH_DIR="${DATA_DIR}/tpch_sf1"
+ RESULTS_FILE="${RESULTS_DIR}/external_aggr.json"
+ echo "RESULTS_FILE: ${RESULTS_FILE}"
+ echo "Running external aggregation benchmark..."
+
+ # Only parquet is supported.
+ # Since per-operator memory limit is calculated as (total-memory-limit /
+ # number-of-partitions), and by default `--partitions` is set to number of
+ # CPU cores, we set a constant number of partitions to prevent this
+ # benchmark to fail on some machines.
+ $CARGO_COMMAND --bin external_aggr -- benchmark --partitions 4
--iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}"
+}
compare_benchmarks() {
diff --git a/benchmarks/src/bin/external_aggr.rs
b/benchmarks/src/bin/external_aggr.rs
new file mode 100644
index 0000000000..1bc74e22cc
--- /dev/null
+++ b/benchmarks/src/bin/external_aggr.rs
@@ -0,0 +1,390 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! external_aggr binary entrypoint
+
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::sync::OnceLock;
+use structopt::StructOpt;
+
+use arrow::record_batch::RecordBatch;
+use arrow::util::pretty;
+use datafusion::datasource::file_format::parquet::ParquetFormat;
+use datafusion::datasource::listing::{
+ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
+};
+use datafusion::datasource::{MemTable, TableProvider};
+use datafusion::error::Result;
+use datafusion::execution::memory_pool::FairSpillPool;
+use datafusion::execution::memory_pool::{human_readable_size, units};
+use datafusion::execution::runtime_env::RuntimeConfig;
+use datafusion::physical_plan::display::DisplayableExecutionPlan;
+use datafusion::physical_plan::{collect, displayable};
+use datafusion::prelude::*;
+use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt};
+use datafusion_common::instant::Instant;
+use datafusion_common::{exec_datafusion_err, exec_err,
DEFAULT_PARQUET_EXTENSION};
+
+#[derive(Debug, StructOpt)]
+#[structopt(
+ name = "datafusion-external-aggregation",
+ about = "DataFusion external aggregation benchmark"
+)]
+enum ExternalAggrOpt {
+ Benchmark(ExternalAggrConfig),
+}
+
+#[derive(Debug, StructOpt)]
+struct ExternalAggrConfig {
+ /// Query number. If not specified, runs all queries
+ #[structopt(short, long)]
+ query: Option<usize>,
+
+ /// Memory limit (e.g. '100M', '1.5G'). If not specified, run all
pre-defined memory limits for given query.
+ #[structopt(long)]
+ memory_limit: Option<String>,
+
+ /// Common options
+ #[structopt(flatten)]
+ common: CommonOpt,
+
+ /// Path to data files (lineitem). Only parquet format is supported
+ #[structopt(parse(from_os_str), required = true, short = "p", long =
"path")]
+ path: PathBuf,
+
+ /// Load the data into a MemTable before executing the query
+ #[structopt(short = "m", long = "mem-table")]
+ mem_table: bool,
+
+ /// Path to JSON benchmark result to be compare using `compare.py`
+ #[structopt(parse(from_os_str), short = "o", long = "output")]
+ output_path: Option<PathBuf>,
+}
+
+struct QueryResult {
+ elapsed: std::time::Duration,
+ row_count: usize,
+}
+
+/// Query Memory Limits
+/// Map query id to predefined memory limits
+///
+/// Q1 requires 36MiB for aggregation
+/// Memory limits to run: 64MiB, 32MiB, 16MiB
+/// Q2 requires 250MiB for aggregation
+/// Memory limits to run: 512MiB, 256MiB, 128MiB, 64MiB, 32MiB
+static QUERY_MEMORY_LIMITS: OnceLock<HashMap<usize, Vec<u64>>> =
OnceLock::new();
+
+impl ExternalAggrConfig {
+ const AGGR_TABLES: [&'static str; 1] = ["lineitem"];
+ const AGGR_QUERIES: [&'static str; 2] = [
+ // Q1: Output size is ~25% of lineitem table
+ r#"
+ SELECT count(*)
+ FROM (
+ SELECT DISTINCT l_orderkey
+ FROM lineitem
+ )
+ "#,
+ // Q2: Output size is ~99% of lineitem table
+ r#"
+ SELECT count(*)
+ FROM (
+ SELECT DISTINCT l_orderkey, l_suppkey
+ FROM lineitem
+ )
+ "#,
+ ];
+
+ fn init_query_memory_limits() -> &'static HashMap<usize, Vec<u64>> {
+ use units::*;
+ QUERY_MEMORY_LIMITS.get_or_init(|| {
+ let mut map = HashMap::new();
+ map.insert(1, vec![64 * MB, 32 * MB, 16 * MB]);
+ map.insert(2, vec![512 * MB, 256 * MB, 128 * MB, 64 * MB, 32 *
MB]);
+ map
+ })
+ }
+
+ /// If `--query` and `--memory-limit` is not speicified, run all queries
+ /// with pre-configured memory limits
+ /// If only `--query` is specified, run the query with all memory limits
+ /// for this query
+ /// If both `--query` and `--memory-limit` are specified, run the query
+ /// with the specified memory limit
+ pub async fn run(&self) -> Result<()> {
+ let mut benchmark_run = BenchmarkRun::new();
+
+ let memory_limit = match &self.memory_limit {
+ Some(limit) => Some(Self::parse_memory_limit(limit)?),
+ None => None,
+ };
+
+ let query_range = match self.query {
+ Some(query_id) => query_id..=query_id,
+ None => 1..=Self::AGGR_QUERIES.len(),
+ };
+
+ // Each element is (query_id, memory_limit)
+ // e.g. [(1, 64_000), (1, 32_000)...] means first run Q1 with 64KiB
+ // memory limit, next run Q1 with 32KiB memory limit, etc.
+ let mut query_executions = vec![];
+ // Setup `query_executions`
+ for query_id in query_range {
+ if query_id > Self::AGGR_QUERIES.len() {
+ return exec_err!(
+ "Invalid '--query'(query number) {} for external
aggregation benchmark.",
+ query_id
+ );
+ }
+
+ match memory_limit {
+ Some(limit) => {
+ query_executions.push((query_id, limit));
+ }
+ None => {
+ let memory_limits_table = Self::init_query_memory_limits();
+ let memory_limits =
memory_limits_table.get(&query_id).unwrap();
+ for limit in memory_limits {
+ query_executions.push((query_id, *limit));
+ }
+ }
+ }
+ }
+
+ for (query_id, mem_limit) in query_executions {
+ benchmark_run.start_new_case(&format!(
+ "{query_id}({})",
+ human_readable_size(mem_limit as usize)
+ ));
+
+ let query_results = self.benchmark_query(query_id,
mem_limit).await?;
+ for iter in query_results {
+ benchmark_run.write_iter(iter.elapsed, iter.row_count);
+ }
+ }
+
+ benchmark_run.maybe_write_json(self.output_path.as_ref())?;
+
+ Ok(())
+ }
+
+ /// Benchmark query `query_id` in `AGGR_QUERIES`
+ async fn benchmark_query(
+ &self,
+ query_id: usize,
+ mem_limit: u64,
+ ) -> Result<Vec<QueryResult>> {
+ let query_name =
+ format!("Q{query_id}({})", human_readable_size(mem_limit as
usize));
+ let mut config = self.common.config();
+ config
+ .options_mut()
+ .execution
+ .parquet
+ .schema_force_view_types = self.common.force_view_types;
+ let runtime_config = RuntimeConfig::new()
+ .with_memory_pool(Arc::new(FairSpillPool::new(mem_limit as usize)))
+ .build_arc()?;
+ let ctx = SessionContext::new_with_config_rt(config, runtime_config);
+
+ // register tables
+ self.register_tables(&ctx).await?;
+
+ let mut millis = vec![];
+ // run benchmark
+ let mut query_results = vec![];
+ for i in 0..self.iterations() {
+ let start = Instant::now();
+
+ let query_idx = query_id - 1; // 1-indexed -> 0-indexed
+ let sql = Self::AGGR_QUERIES[query_idx];
+
+ let result = self.execute_query(&ctx, sql).await?;
+
+ let elapsed = start.elapsed(); //.as_secs_f64() * 1000.0;
+ let ms = elapsed.as_secs_f64() * 1000.0;
+ millis.push(ms);
+
+ let row_count = result.iter().map(|b| b.num_rows()).sum();
+ println!(
+ "{query_name} iteration {i} took {ms:.1} ms and returned
{row_count} rows"
+ );
+ query_results.push(QueryResult { elapsed, row_count });
+ }
+
+ let avg = millis.iter().sum::<f64>() / millis.len() as f64;
+ println!("{query_name} avg time: {avg:.2} ms");
+
+ Ok(query_results)
+ }
+
+ async fn register_tables(&self, ctx: &SessionContext) -> Result<()> {
+ for table in Self::AGGR_TABLES {
+ let table_provider = { self.get_table(ctx, table).await? };
+
+ if self.mem_table {
+ println!("Loading table '{table}' into memory");
+ let start = Instant::now();
+ let memtable =
+ MemTable::load(table_provider, Some(self.partitions()),
&ctx.state())
+ .await?;
+ println!(
+ "Loaded table '{}' into memory in {} ms",
+ table,
+ start.elapsed().as_millis()
+ );
+ ctx.register_table(table, Arc::new(memtable))?;
+ } else {
+ ctx.register_table(table, table_provider)?;
+ }
+ }
+ Ok(())
+ }
+
+ async fn execute_query(
+ &self,
+ ctx: &SessionContext,
+ sql: &str,
+ ) -> Result<Vec<RecordBatch>> {
+ let debug = self.common.debug;
+ let plan = ctx.sql(sql).await?;
+ let (state, plan) = plan.into_parts();
+
+ if debug {
+ println!("=== Logical plan ===\n{plan}\n");
+ }
+
+ let plan = state.optimize(&plan)?;
+ if debug {
+ println!("=== Optimized logical plan ===\n{plan}\n");
+ }
+ let physical_plan = state.create_physical_plan(&plan).await?;
+ if debug {
+ println!(
+ "=== Physical plan ===\n{}\n",
+ displayable(physical_plan.as_ref()).indent(true)
+ );
+ }
+ let result = collect(physical_plan.clone(), state.task_ctx()).await?;
+ if debug {
+ println!(
+ "=== Physical plan with metrics ===\n{}\n",
+ DisplayableExecutionPlan::with_metrics(physical_plan.as_ref())
+ .indent(true)
+ );
+ if !result.is_empty() {
+ // do not call print_batches if there are no batches as the
result is confusing
+ // and makes it look like there is a batch with no columns
+ pretty::print_batches(&result)?;
+ }
+ }
+ Ok(result)
+ }
+
+ async fn get_table(
+ &self,
+ ctx: &SessionContext,
+ table: &str,
+ ) -> Result<Arc<dyn TableProvider>> {
+ let path = self.path.to_str().unwrap();
+
+ // Obtain a snapshot of the SessionState
+ let state = ctx.state();
+ let path = format!("{path}/{table}");
+ let format = Arc::new(
+ ParquetFormat::default()
+ .with_options(ctx.state().table_options().parquet.clone()),
+ );
+ let extension = DEFAULT_PARQUET_EXTENSION;
+
+ let options = ListingOptions::new(format)
+ .with_file_extension(extension)
+ .with_collect_stat(state.config().collect_statistics());
+
+ let table_path = ListingTableUrl::parse(path)?;
+ let config =
ListingTableConfig::new(table_path).with_listing_options(options);
+ let config = config.infer_schema(&state).await?;
+
+ Ok(Arc::new(ListingTable::try_new(config)?))
+ }
+
+ fn iterations(&self) -> usize {
+ self.common.iterations
+ }
+
+ fn partitions(&self) -> usize {
+ self.common.partitions.unwrap_or(num_cpus::get())
+ }
+
+ /// Parse memory limit from string to number of bytes
+ /// e.g. '1.5G', '100M' -> 1572864
+ fn parse_memory_limit(limit: &str) -> Result<u64> {
+ let (number, unit) = limit.split_at(limit.len() - 1);
+ let number: f64 = number.parse().map_err(|_| {
+ exec_datafusion_err!("Failed to parse number from memory limit
'{}'", limit)
+ })?;
+
+ match unit {
+ "K" => Ok((number * 1024.0) as u64),
+ "M" => Ok((number * 1024.0 * 1024.0) as u64),
+ "G" => Ok((number * 1024.0 * 1024.0 * 1024.0) as u64),
+ _ => exec_err!("Unsupported unit '{}' in memory limit '{}'", unit,
limit),
+ }
+ }
+}
+
+#[tokio::main]
+pub async fn main() -> Result<()> {
+ env_logger::init();
+
+ match ExternalAggrOpt::from_args() {
+ ExternalAggrOpt::Benchmark(opt) => opt.run().await?,
+ }
+
+ Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_parse_memory_limit_all() {
+ // Test valid inputs
+ assert_eq!(
+ ExternalAggrConfig::parse_memory_limit("100K").unwrap(),
+ 102400
+ );
+ assert_eq!(
+ ExternalAggrConfig::parse_memory_limit("1.5M").unwrap(),
+ 1572864
+ );
+ assert_eq!(
+ ExternalAggrConfig::parse_memory_limit("2G").unwrap(),
+ 2147483648
+ );
+
+ // Test invalid unit
+ assert!(ExternalAggrConfig::parse_memory_limit("500X").is_err());
+
+ // Test invalid number
+ assert!(ExternalAggrConfig::parse_memory_limit("abcM").is_err());
+ }
+}
diff --git a/datafusion/execution/src/memory_pool/mod.rs
b/datafusion/execution/src/memory_pool/mod.rs
index d87ce1ebfe..5bf30b724d 100644
--- a/datafusion/execution/src/memory_pool/mod.rs
+++ b/datafusion/execution/src/memory_pool/mod.rs
@@ -334,13 +334,17 @@ impl Drop for MemoryReservation {
}
}
-const TB: u64 = 1 << 40;
-const GB: u64 = 1 << 30;
-const MB: u64 = 1 << 20;
-const KB: u64 = 1 << 10;
+pub mod units {
+ pub const TB: u64 = 1 << 40;
+ pub const GB: u64 = 1 << 30;
+ pub const MB: u64 = 1 << 20;
+ pub const KB: u64 = 1 << 10;
+}
/// Present size in human readable form
pub fn human_readable_size(size: usize) -> String {
+ use units::*;
+
let size = size as u64;
let (value, unit) = {
if size >= 2 * TB {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
