Skip to content

Commit ec64cba

Browse files
committed
Add adaptive plan filter execution
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 819ed4a commit ec64cba

12 files changed

Lines changed: 297 additions & 35 deletions

File tree

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

vortex-datafusion/src/persistent/opener.rs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ use vortex::metrics::Label;
5454
use vortex::metrics::MetricsRegistry;
5555
use vortex::session::VortexSession;
5656
use vortex_arrow::ArrowSessionExt;
57+
use vortex_scan_v2::FilterMode;
5758
use vortex_scan_v2::ScanBuilder as PlanScanBuilder;
5859
use vortex_utils::aliases::dash_map::DashMap;
5960
use vortex_utils::aliases::dash_map::Entry;
@@ -386,6 +387,13 @@ impl FileOpener for VortexOpener {
386387
}
387388
}
388389

390+
let filter_mode = if std::env::var("VORTEX_PLAN_V2_FILTER_MODE").as_deref()
391+
== Ok("adaptive")
392+
{
393+
FilterMode::Adaptive
394+
} else {
395+
FilterMode::Parallel
396+
};
389397
let mut scan_builder = PlanScanBuilder::try_new(
390398
vxf.footer().layout(),
391399
vxf.segment_source(),
@@ -402,7 +410,8 @@ impl FileOpener for VortexOpener {
402410
.map(unbind)
403411
.transpose()
404412
.map_err(|error| DataFusionError::External(Box::new(error)))?,
405-
);
413+
)
414+
.with_filter_mode(filter_mode);
406415
if let Some(limit) = limit
407416
&& filter.is_none()
408417
{

vortex-layout/src/plan/plans/zoned.rs

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ use vortex_session::registry::CachedId;
3535

3636
use crate::layouts::zoned::zone_map::ZoneMap;
3737
use crate::plan::Eval;
38+
use crate::plan::EvalPlan;
3839
use crate::plan::Plan;
3940
use crate::plan::PlanArrayFuture;
4041
use crate::plan::PlanChildren;
@@ -223,6 +224,26 @@ impl ZonedPlan {
223224
))
224225
}
225226

227+
fn with_data_expression(&self, expression: BoundExpression) -> VortexResult<Option<Self>> {
228+
let Some(data_plan) = self.data_plan()? else {
229+
return Ok(None);
230+
};
231+
Ok(Some(
232+
PlanParts {
233+
vtable: Zoned,
234+
dtype: expression.dtype().clone(),
235+
row_count: self.row_count(),
236+
children: vec![
237+
EvalPlan::try_new(expression, data_plan)?.into_plan(),
238+
self.zones_plan()?,
239+
]
240+
.into(),
241+
data: self.data().clone(),
242+
}
243+
.into_typed(),
244+
))
245+
}
246+
226247
fn execute_pruning(
227248
&self,
228249
ctx: &PlanExecutionContext,
@@ -380,7 +401,6 @@ impl PlanVTable for Zoned {
380401
if data_plan.dtype() != plan.dtype() || data_plan.row_count() != plan.row_count() {
381402
vortex_error::vortex_bail!("Zoned data child shape does not match the plan output");
382403
}
383-
data.column_dtype = plan.dtype().clone();
384404
Ok(())
385405
}
386406

@@ -414,7 +434,7 @@ impl PlanVTable for Zoned {
414434
}
415435
}
416436

417-
/// Rewrites an abstract statistic expression over a zoned plan into its pruning state.
437+
/// Pushes data expressions through a zoned plan and rewrites statistic expressions into pruning.
418438
#[derive(Debug)]
419439
pub(crate) struct ExpressionZonedRule;
420440

@@ -444,12 +464,17 @@ impl PlanParentReduceRule<Zoned> for ExpressionZonedRule {
444464
contains_root |= expression.is_root();
445465
Ok(Transformed::no(expression))
446466
})?;
447-
if !parent.dtype().is_boolean() || !contains_stat || contains_root {
448-
return Ok(None);
467+
if contains_stat {
468+
if !parent.dtype().is_boolean() || contains_root {
469+
return Ok(None);
470+
}
471+
return Ok(child
472+
.with_pruning(parent.expression().clone())?
473+
.map(Plan::into_plan));
449474
}
450475

451476
Ok(child
452-
.with_pruning(parent.expression().clone())?
477+
.with_data_expression(parent.expression().clone())?
453478
.map(Plan::into_plan))
454479
}
455480
}

vortex-layout/src/plan/tests.rs

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1003,6 +1003,37 @@ fn zoned_plan_exposes_data_and_zones() -> VortexResult<()> {
10031003
Ok(())
10041004
}
10051005

1006+
#[test]
1007+
fn data_expression_pushes_through_zoned_and_preserves_zones() -> VortexResult<()> {
1008+
let dtype = primitive(PType::I32, Nullability::NonNullable);
1009+
let zones_dtype = DType::Struct(StructFields::empty(), Nullability::NonNullable);
1010+
let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?;
1011+
let aggregate_fns: Arc<[AggregateFnRef]> = Vec::new().into();
1012+
let layout = ZonedLayout::try_new(
1013+
flat(5, dtype.clone(), 0),
1014+
flat(2, zones_dtype, 1),
1015+
zone_len,
1016+
aggregate_fns,
1017+
)?
1018+
.into_layout();
1019+
let source = make_plan(layout)?;
1020+
let expression = gt(root(), lit(5_i32)).bind(&dtype)?;
1021+
1022+
let optimized = optimize(EvalPlan::try_new(expression, source)?.into_plan())?;
1023+
insta::assert_snapshot!(optimized.display_tree(), @r"
1024+
root: vortex.plan.zoned(bool, rows=5)
1025+
data: vortex.plan.eval(bool, rows=5) expr=($ > 5i32)
1026+
child: vortex.plan.segment_scan(i32, rows=5)
1027+
zones: vortex.plan.segment_scan({}, rows=2)
1028+
");
1029+
let optimized_zoned = optimized
1030+
.as_opt::<Zoned>()
1031+
.ok_or_else(|| vortex_err!("optimized plan is not zoned"))?;
1032+
assert!(!optimized_zoned.is_pruning());
1033+
assert_eq!(optimized_zoned.zones_plan()?.row_count(), 2);
1034+
Ok(())
1035+
}
1036+
10061037
#[test]
10071038
fn stats_expression_rewrites_to_zoned_pruning_plan() -> VortexResult<()> {
10081039
let dtype = primitive(PType::I32, Nullability::NonNullable);

vortex-layout/src/scan/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ mod tasks;
1313
#[cfg(test)]
1414
mod test;
1515

16+
pub use filter::FilterExpr;
17+
1618
/// A heuristic for an ideal split size.
1719
///
1820
/// We don't actually know if this is right, but it is probably a good estimate.

vortex-scan-v2/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ rust-version = { workspace = true }
1414
version = { workspace = true }
1515

1616
[dependencies]
17+
bit-vec = { workspace = true }
1718
futures = { workspace = true, features = ["alloc", "async-await"] }
1819
itertools = { workspace = true }
1920
tracing = { workspace = true }

vortex-scan-v2/src/filter.rs

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use std::ops::Range;
5+
use std::sync::Arc;
6+
7+
use bit_vec::BitVec;
8+
use vortex_array::MaskFuture;
9+
use vortex_array::VortexSessionExecute;
10+
use vortex_error::VortexResult;
11+
use vortex_layout::plan::PlanExecutionContext;
12+
use vortex_layout::plan::PlanRef;
13+
use vortex_layout::scan::FilterExpr;
14+
use vortex_mask::Mask;
15+
16+
/// Controls how a scan executes top-level filter conjunctions.
17+
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
18+
pub enum FilterMode {
19+
/// Optimizes the complete predicate as one plan, allowing independent branches to run in
20+
/// parallel.
21+
#[default]
22+
Parallel,
23+
/// Splits top-level conjunctions and executes them as an adaptively ordered mask chain.
24+
Adaptive,
25+
}
26+
27+
#[derive(Clone)]
28+
pub(crate) enum FilterPlan {
29+
Parallel(PlanRef),
30+
Adaptive {
31+
filter: Arc<FilterExpr>,
32+
plans: Arc<[PlanRef]>,
33+
},
34+
}
35+
36+
impl FilterPlan {
37+
pub(crate) fn parallel(plan: PlanRef) -> Self {
38+
Self::Parallel(plan)
39+
}
40+
41+
pub(crate) fn adaptive(filter: FilterExpr, plans: Vec<PlanRef>) -> Self {
42+
Self::Adaptive {
43+
filter: Arc::new(filter),
44+
plans: plans.into(),
45+
}
46+
}
47+
48+
pub(crate) fn plans(&self) -> Vec<&PlanRef> {
49+
match self {
50+
Self::Parallel(plan) => vec![plan],
51+
Self::Adaptive { plans, .. } => plans.iter().collect(),
52+
}
53+
}
54+
55+
pub(crate) fn execute(
56+
&self,
57+
execution: &PlanExecutionContext,
58+
row_range: &Range<u64>,
59+
row_mask: Mask,
60+
) -> VortexResult<MaskFuture> {
61+
match self {
62+
Self::Parallel(filter) => {
63+
let predicate =
64+
filter.execute(execution, row_range, MaskFuture::ready(row_mask.clone()))?;
65+
let session = execution.session().clone();
66+
Ok(MaskFuture::new(row_mask.len(), async move {
67+
let predicate = predicate.await?;
68+
let mut execution = session.create_execution_ctx();
69+
let predicate: Mask = predicate.null_as_false().execute(&mut execution)?;
70+
Ok(row_mask.intersect_by_rank(&predicate))
71+
}))
72+
}
73+
Self::Adaptive { filter, plans } => {
74+
let execution = execution.clone();
75+
let row_range = row_range.clone();
76+
let filter = Arc::clone(filter);
77+
let plans = Arc::clone(plans);
78+
Ok(MaskFuture::new(row_mask.len(), async move {
79+
let mut row_mask = row_mask;
80+
let mut remaining = BitVec::from_elem(plans.len(), true);
81+
while let Some(index) = filter.next_conjunct(&remaining) {
82+
remaining.set(index, false);
83+
if row_mask.all_false() {
84+
break;
85+
}
86+
87+
let input_rows = row_mask.true_count();
88+
let predicate = plans[index].execute(
89+
&execution,
90+
&row_range,
91+
MaskFuture::ready(row_mask.clone()),
92+
)?;
93+
let predicate = predicate.await?;
94+
let mut ctx = execution.session().create_execution_ctx();
95+
let predicate: Mask = predicate.null_as_false().execute(&mut ctx)?;
96+
row_mask = row_mask.intersect_by_rank(&predicate);
97+
filter.report_selectivity(index, row_mask.density());
98+
tracing::trace!(
99+
target: "vortex_scan_v2::execution",
100+
conjunct = index,
101+
input_rows,
102+
output_rows = row_mask.true_count(),
103+
"applied an adaptive plan filter conjunct"
104+
);
105+
}
106+
Ok(row_mask)
107+
}))
108+
}
109+
}
110+
}
111+
}

vortex-scan-v2/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
//! Set `RUST_LOG=vortex_scan_v2=debug` to log source and optimized plan trees and selected scan
1111
//! splits. Use `trace` to also log execution of each split.
1212
13+
mod filter;
1314
mod repeated_scan;
1415
mod scan_builder;
1516
mod splits;
@@ -18,6 +19,7 @@ mod tasks;
1819
#[cfg(test)]
1920
mod tests;
2021

22+
pub use filter::FilterMode;
2123
pub use repeated_scan::RepeatedScan;
2224
pub use scan_builder::ScanBuilder;
2325
pub use splits::SplitBy;

vortex-scan-v2/src/repeated_scan.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use vortex_layout::plan::PlanRef;
2525
use vortex_scan::selection::Selection;
2626
use vortex_utils::parallelism::get_available_parallelism;
2727

28+
use crate::filter::FilterPlan;
2829
use crate::splits::Splits;
2930
use crate::tasks::TaskContext;
3031
use crate::tasks::split_exec;
@@ -34,7 +35,7 @@ pub struct RepeatedScan<A: 'static + Send> {
3435
execution: PlanExecutionContext,
3536
projection: PlanRef,
3637
pruning: Option<PlanRef>,
37-
filter: Option<PlanRef>,
38+
filter: Option<FilterPlan>,
3839
ordered: bool,
3940
row_range: Option<Range<u64>>,
4041
selection: Selection,
@@ -82,7 +83,7 @@ impl<A: 'static + Send> RepeatedScan<A> {
8283
execution: PlanExecutionContext,
8384
projection: PlanRef,
8485
pruning: Option<PlanRef>,
85-
filter: Option<PlanRef>,
86+
filter: Option<FilterPlan>,
8687
ordered: bool,
8788
row_range: Option<Range<u64>>,
8889
selection: Selection,

0 commit comments

Comments
 (0)