Skip to content

feat: add bounded Parquet row-group prefetch #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
peterxcli wants to merge 1 commit into from
Draft

feat: add bounded Parquet row-group prefetch #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 49 additions & 1 deletion benchmarks/src/clickbench.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,16 +18,21 @@
use std::fs;
use std::io::ErrorKind;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats};
use clap::Args;
use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource};
use datafusion::datasource::source::DataSourceExec;
use datafusion::logical_expr::{ExplainFormat, ExplainOption};
use datafusion::physical_plan::collect;
use datafusion::{
error::{DataFusionError, Result},
prelude::SessionContext,
};
use datafusion_common::exec_datafusion_err;
use datafusion_common::instant::Instant;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};

/// SQL to create the hits view with proper EventDate casting.
///
Expand Down Expand Up @@ -62,6 +67,11 @@ pub struct RunOpt {
#[arg(long = "pushdown")]
pushdown: bool,

/// Maximum compressed bytes to prefetch for the next Parquet row group.
/// Zero disables prefetch. Uses the query's memory pool.
#[arg(long, default_value_t = 0)]
prefetch_bytes: usize,

/// Common options
#[command(flatten)]
common: CommonOpt,
Expand Down Expand Up @@ -254,8 +264,46 @@ impl RunOpt {
let mut query_results = vec![];
for i in 0..self.iterations() {
let start = Instant::now();
let results = ctx.sql(sql).await?.collect().await?;
let dataframe = ctx.sql(sql).await?;
let task_ctx = Arc::new(dataframe.task_ctx());
let mut plan = dataframe.create_physical_plan().await?;
if self.prefetch_bytes > 0 {
plan = plan
.transform_up(|plan| {
if let Some(exec) = plan.downcast_ref::<DataSourceExec>()
&& let Some((config, source)) =
exec.downcast_to_file_source::<ParquetSource>()
{
let source = source.clone().with_row_group_prefetch(
self.prefetch_bytes,
Arc::clone(&ctx.runtime_env().memory_pool),
);
let config = FileScanConfigBuilder::from(config.clone())
.with_source(Arc::new(source))
.build();
return Ok(Transformed::yes(Arc::new(
exec.clone().with_data_source(Arc::new(config)),
)));
}
Ok(Transformed::no(plan))
})?
.data;
}
let results = collect(Arc::clone(&plan), task_ctx).await?;
let elapsed = start.elapsed();
if self.prefetch_bytes > 0 {
let mut prefetched = 0;
plan.apply(|plan| {
if let Some(value) = plan
.metrics()
.and_then(|metrics| metrics.sum_by_name("prefetch_row_groups"))
{
prefetched += value.as_usize();
}
Ok(TreeNodeRecursion::Continue)
})?;
println!("Prefetched row groups: {prefetched}");
}
let ms = elapsed.as_secs_f64() * 1000.0;
millis.push(ms);
let row_count: usize = results.iter().map(|b| b.num_rows()).sum();
Expand Down
168 changes: 96 additions & 72 deletions datafusion/core/tests/parquet/filter_pushdown.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -814,9 +814,11 @@ async fn pushed_down_predicate_reports_the_original_error() {
async fn upfront_io_preserves_page_pruning_from_session_configuration() {
use arrow::array::{Int64Array, StringArray, StructArray};
use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
use datafusion_datasource::source::DataSourceExec;
use datafusion_datasource_parquet::source::ParquetSource;
use datafusion_execution::memory_pool::GreedyMemoryPool;
use parquet::arrow::ArrowWriter;
use parquet::file::properties::EnabledStatistics;
use std::sync::Arc;
Expand Down Expand Up @@ -904,82 +906,104 @@ async fn upfront_io_preserves_page_pruning_from_session_configuration() {
for (pushdown, progressive) in
[(false, true), (true, true), (true, false)]
{
let config = SessionConfig::new()
.with_target_partitions(1)
.with_batch_size(127);
let ctx = SessionContext::new_with_config(config);
for (key, value) in [
("pushdown_filters", pushdown),
("progressive_io", progressive),
] {
ctx.sql(&format!(
"SET datafusion.execution.parquet.{key} = {value}"
))
.await
.unwrap()
.collect()
for budget in [0, 1 << 20] {
let config = SessionConfig::new()
.with_target_partitions(1)
.with_batch_size(127);
let ctx = SessionContext::new_with_config(config);
for (key, value) in [
("pushdown_filters", pushdown),
("progressive_io", progressive),
] {
ctx.sql(&format!(
"SET datafusion.execution.parquet.{key} = {value}"
))
.await
.unwrap()
.collect()
.await
.unwrap();
}
ctx.register_parquet(
"t",
path.to_str().unwrap(),
ParquetReadOptions::default(),
)
.await
.unwrap();
}
ctx.register_parquet(
"t",
path.to_str().unwrap(),
ParquetReadOptions::default(),
)
.await
.unwrap();
let key = if predicate_in_output { "key," } else { "" };
let plan = ctx.sql(&format!("SELECT {key} payload, nested.value AS value FROM t WHERE key < 100 ORDER BY value"))
.await.unwrap().create_physical_plan().await.unwrap();
let mut scans = 0;
plan.apply(|plan| {
if let Some(exec) = plan.downcast_ref::<DataSourceExec>()
&& let Some((_, source)) =
exec.downcast_to_file_source::<ParquetSource>()
{
scans += 1;
assert_eq!(
source.table_parquet_options().global.progressive_io,
progressive
);
let key = if predicate_in_output { "key," } else { "" };
let plan = ctx.sql(&format!("SELECT {key} payload, nested.value AS value FROM t WHERE key < 100 ORDER BY value"))
.await.unwrap().create_physical_plan().await.unwrap();
let mut scans = 0;
let plan = plan
.transform_up(|plan| {
if let Some(exec) = plan.downcast_ref::<DataSourceExec>()
&& let Some((config, source)) =
exec.downcast_to_file_source::<ParquetSource>()
{
scans += 1;
assert_eq!(
source
.table_parquet_options()
.global
.progressive_io,
progressive
);
let source = source.clone().with_row_group_prefetch(
budget,
Arc::new(GreedyMemoryPool::new(1 << 20)),
);
let config =
FileScanConfigBuilder::from(config.clone())
.with_source(Arc::new(source))
.build();
return Ok(Transformed::yes(Arc::new(
exec.clone().with_data_source(Arc::new(config)),
)));
}
Ok(Transformed::no(plan))
})
.unwrap()
.data;
assert_eq!(scans, 1);
let batches =
collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
let batch = concat_batches(&plan.schema(), &batches).unwrap();
let values = batch
.column(batch.num_columns() - 1)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
assert_eq!(values.values().as_ref(), expected.as_slice());
if let Some(expected) = &expected_batch {
assert_eq!(&batch, expected);
} else {
expected_batch = Some(batch);
}
Ok(TreeNodeRecursion::Continue)
})
.unwrap();
assert_eq!(scans, 1);
let batches =
collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
let batch = concat_batches(&plan.schema(), &batches).unwrap();
let values = batch
.column(batch.num_columns() - 1)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
assert_eq!(values.values().as_ref(), expected.as_slice());
if let Some(expected) = &expected_batch {
assert_eq!(&batch, expected);
} else {
expected_batch = Some(batch);
}
let metrics = TestParquetFile::parquet_metrics(&plan).unwrap();
let bytes = get_value(&metrics, "bytes_scanned");
if indexed && clustered {
assert!(
get_pruning_metrics(&metrics, "page_index_rows_pruned").0 > 0
);
}
if !predicate_in_output {
if let Some(expected) = baseline_bytes {
assert_eq!(
bytes, expected,
"clustered={clustered}, indexed={indexed}, pushdown={pushdown}, progressive={progressive}"
let metrics = TestParquetFile::parquet_metrics(&plan).unwrap();
let bytes = get_value(&metrics, "bytes_scanned");
if indexed && clustered {
assert!(
get_pruning_metrics(&metrics, "page_index_rows_pruned").0
> 0
);
} else {
baseline_bytes = Some(bytes);
}
}
if pushdown {
assert!(get_value(&metrics, "pushdown_rows_pruned") > 0);
if !predicate_in_output {
if let Some(expected) = baseline_bytes {
assert_eq!(
bytes, expected,
"clustered={clustered}, indexed={indexed}, pushdown={pushdown}, progressive={progressive}, budget={budget}"
);
} else {
baseline_bytes = Some(bytes);
}
}
if budget > 0 {
assert!(get_value(&metrics, "prefetch_bytes") > 0);
}
if pushdown {
assert!(get_value(&metrics, "pushdown_rows_pruned") > 0);
}
}
}
}
Expand Down
1 change: 1 addition & 0 deletions datafusion/datasource-parquet/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ criterion = { workspace = true }
datafusion-functions = { workspace = true }
datafusion-functions-nested = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["rt-multi-thread"] }

# Note: add additional linter rules in lib.rs.
# Rust does not support workspace + new linter rules in subcrates yet
Expand Down
23 changes: 23 additions & 0 deletions datafusion/datasource-parquet/src/metrics.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,3 +467,26 @@ impl RowFilterSkippedFullyMatchedMetric {
count.add(1);
}
}

/// Per-stream counters registered in the scan's existing metric set.
#[derive(Clone)]
pub(crate) struct PrefetchMetrics {
pub bytes: Count,
pub row_groups: Count,
pub budget_skips: Count,
pub wait_time: Time,
}

impl PrefetchMetrics {
pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self {
Self {
bytes: MetricBuilder::new(metrics).counter("prefetch_bytes", partition),
row_groups: MetricBuilder::new(metrics)
.counter("prefetch_row_groups", partition),
budget_skips: MetricBuilder::new(metrics)
.counter("prefetch_budget_skips", partition),
wait_time: MetricBuilder::new(metrics)
.subset_time("prefetch_wait_time", partition),
}
}
}
Loading
Loading