Skip to content

Sink a filter that holds a projection into NdFilterExec (#397) #401

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

Open
robinskil wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Sink a filter that holds a projection into NdFilterExec (#397) #401

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
204 changes: 204 additions & 0 deletions beacon-db/beacon-datafusion-ext/src/nd/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -856,4 +856,208 @@ mod tests {
"expected FilterExec → NdBroadcastExec → NdFilterExec:\n{rendered}"
);
}
/// A narrow select list gives the filter a projection (`FilterExec: …,
/// projection=[…]`). The projection is a plain column list. So it sinks with
/// the predicate. An `NdProjectionExec` goes between the nd filter and the
/// broadcast. No `FilterExec` stays above.
#[tokio::test]
async fn pushdown_rule_sinks_a_narrowing_filter_projection() {
use datafusion::common::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::displayable;
use datafusion::physical_plan::filter::FilterExecBuilder;

let schema = test_source().schema();
// The predicate reads `lon`. The select list keeps only `lat`.
let predicate: Arc<dyn PhysicalExpr> =
binary(col("lon", &schema).unwrap(), Operator::Eq, lit(5i32), &schema).unwrap();

let original: Arc<dyn ExecutionPlan> = Arc::new(
FilterExecBuilder::new(
predicate,
Arc::new(NdBroadcastExec::try_new(test_source()).unwrap()),
)
.apply_projection(Some(vec![schema.index_of("lat").unwrap()]))
.unwrap()
.build()
.unwrap(),
);
let original_schema = original.schema();
let expected = run(original.clone()).await.unwrap();

let optimized = NdFilterPushdown::new()
.optimize(original, &ConfigOptions::default())
.unwrap();

assert_eq!(optimized.schema(), original_schema);

let rendered = displayable(optimized.as_ref()).indent(true).to_string();
assert!(
!rendered
.lines()
.any(|l| l.trim_start().starts_with("FilterExec:")),
"the full filter sinks, so no FilterExec stays above:\n{rendered}"
);
let broadcast = rendered.find("NdBroadcastExec");
let projection = rendered.find("NdProjectionExec");
let filter = rendered.find("NdFilterExec");
let source = rendered.find("NdSourceExec");
assert!(
broadcast < projection && projection < filter && filter < source,
"expected NdBroadcastExec → NdProjectionExec → NdFilterExec → NdSourceExec:\n{rendered}"
);

let actual = run(optimized).await.unwrap();
assert_eq!(actual, expected);
// One of the two lon values survives. That is half of the 24-cell grid.
assert_eq!(actual.num_rows(), 12);
assert_eq!(actual.num_columns(), 1);
}

/// `count(*)` gives the filter an *empty* projection. The projection sinks
/// the same way. The nd projection keeps the grid selection of the nd filter.
/// So the broadcast reports the cells that remain as rows over no columns.
#[tokio::test]
async fn pushdown_rule_sinks_an_empty_filter_projection() {
use datafusion::common::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::displayable;
use datafusion::physical_plan::filter::FilterExecBuilder;

let schema = test_source().schema();
let predicate: Arc<dyn PhysicalExpr> =
binary(col("lon", &schema).unwrap(), Operator::Eq, lit(5i32), &schema).unwrap();

let original: Arc<dyn ExecutionPlan> = Arc::new(
FilterExecBuilder::new(
predicate,
Arc::new(NdBroadcastExec::try_new(test_source()).unwrap()),
)
.apply_projection(Some(vec![]))
.unwrap()
.build()
.unwrap(),
);
let expected = run(original.clone()).await.unwrap();

let optimized = NdFilterPushdown::new()
.optimize(original, &ConfigOptions::default())
.unwrap();

let rendered = displayable(optimized.as_ref()).indent(true).to_string();
assert!(
rendered.contains("NdFilterExec") && rendered.contains("NdProjectionExec: exprs=[]"),
"the empty projection sinks with the predicate:\n{rendered}"
);

let actual = run(optimized).await.unwrap();
assert_eq!(actual, expected);
// An aggregate counts these rows. The batch holds no column.
assert_eq!(actual.num_rows(), 12);
assert_eq!(actual.num_columns(), 0);
}

/// A conjunct that stays above the broadcast reads the columns that the
/// projection drops. So the projection stays with the residual filter. The
/// element-wise conjunct still sinks.
#[tokio::test]
async fn pushdown_rule_keeps_the_projection_with_a_residual_conjunct() {
use datafusion::common::config::ConfigOptions;
use datafusion::physical_expr::expressions::in_list;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::displayable;
use datafusion::physical_plan::filter::FilterExecBuilder;

let schema = test_source().schema();
// `lon = 5` sinks. `time IN (100, 101)` is outside the whitelist, so it
// stays. Both read a column that the select list drops.
let predicate: Arc<dyn PhysicalExpr> = binary(
binary(col("lon", &schema).unwrap(), Operator::Eq, lit(5i32), &schema).unwrap(),
Operator::And,
in_list(
col("time", &schema).unwrap(),
vec![lit(100i32), lit(101i32)],
&false,
&schema,
)
.unwrap(),
&schema,
)
.unwrap();

let original: Arc<dyn ExecutionPlan> = Arc::new(
FilterExecBuilder::new(
predicate,
Arc::new(NdBroadcastExec::try_new(test_source()).unwrap()),
)
.apply_projection(Some(vec![schema.index_of("lat").unwrap()]))
.unwrap()
.build()
.unwrap(),
);
let original_schema = original.schema();
let expected = run(original.clone()).await.unwrap();

let optimized = NdFilterPushdown::new()
.optimize(original, &ConfigOptions::default())
.unwrap();

assert_eq!(optimized.schema(), original_schema);

let rendered = displayable(optimized.as_ref()).indent(true).to_string();
let residual = rendered
.lines()
.find(|l| l.trim_start().starts_with("FilterExec:"))
.unwrap_or_else(|| panic!("expected a residual FilterExec:\n{rendered}"));
assert!(
residual.contains("projection=["),
"the residual filter keeps the projection:\n{rendered}"
);
assert!(
rendered.contains("NdFilterExec: predicate=[lon@2 = 5]"),
"the element-wise conjunct still sinks:\n{rendered}"
);

let actual = run(optimized).await.unwrap();
assert_eq!(actual, expected);
// The first chunk holds two time steps. At lon = 5 that gives 6 rows.
assert_eq!(actual.num_rows(), 6);
assert_eq!(actual.num_columns(), 1);
}

/// A `fetch` caps the rows that the filter returns. The nd filter holds no
/// cap. So the rule keeps the filter in place, and the cap stays.
#[tokio::test]
async fn pushdown_rule_skips_a_filter_with_a_fetch() {
use datafusion::common::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::displayable;
use datafusion::physical_plan::filter::FilterExecBuilder;

let schema = test_source().schema();
let predicate: Arc<dyn PhysicalExpr> =
binary(col("lon", &schema).unwrap(), Operator::Eq, lit(5i32), &schema).unwrap();

let original: Arc<dyn ExecutionPlan> = Arc::new(
FilterExecBuilder::new(
predicate,
Arc::new(NdBroadcastExec::try_new(test_source()).unwrap()),
)
.with_fetch(Some(3))
.build()
.unwrap(),
);

let optimized = NdFilterPushdown::new()
.optimize(original, &ConfigOptions::default())
.unwrap();

let rendered = displayable(optimized.as_ref()).indent(true).to_string();
assert!(
!rendered.contains("NdFilterExec"),
"a capped filter must stay above the broadcast:\n{rendered}"
);
// The cap still holds.
assert_eq!(run(optimized).await.unwrap().num_rows(), 3);
}
}
86 changes: 73 additions & 13 deletions beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

use std::sync::Arc;

use arrow::datatypes::SchemaRef;
use datafusion::common::config::ConfigOptions;
use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::error::Result;
Expand All @@ -30,7 +31,7 @@ use datafusion::physical_expr::expressions::{
use datafusion::physical_expr::{ScalarFunctionExpr, conjunction, split_conjunction};
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::filter::FilterExec;
use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder};
use datafusion::physical_plan::projection::ProjectionExec;
use datafusion::logical_expr::Volatility;

Expand Down Expand Up @@ -123,8 +124,23 @@ impl PhysicalOptimizerRule for NdProjectionPushdown {
///
/// where `a`, `b` are element-wise ([`is_pushable_expr`]) and `c` is not (e.g. a
/// volatile function or a subquery). If every conjunct is pushable, the residual
/// `FilterExec` is dropped entirely. The rewrite is schema-preserving: a filter
/// never changes columns.
/// `FilterExec` is dropped entirely.
///
/// A `FilterExec` also holds a projection. DataFusion adds one when the select
/// list is narrower than the predicate. `count(*)` gets an empty projection. A
/// filter projection is a plain column list, so it sinks too. It becomes an
/// [`NdProjectionExec`] between the nd filter and the broadcast:
///
/// ```text
/// FilterExec[a, projection=[lat]] NdBroadcastExec
/// NdBroadcastExec -> NdProjectionExec[lat]
/// nd-child NdFilterExec[a]
/// nd-child
/// ```
///
/// A residual conjunct reads the columns that the projection drops. So the
/// projection stays with the residual filter above the broadcast. Both forms keep
/// the schema of the original filter.
#[derive(Debug, Default)]
pub struct NdFilterPushdown;

Expand All @@ -144,15 +160,15 @@ impl PhysicalOptimizerRule for NdFilterPushdown {
let Some(filter) = node.as_any().downcast_ref::<FilterExec>() else {
return Ok(Transformed::no(node));
};
// A `FilterExec` carrying an embedded projection also changes the
// schema; leave those in place so the rewrite stays a pure row
// selection.
if filter.projection().is_some() {
return Ok(Transformed::no(node));
}
let Some(broadcast) = filter.input().as_any().downcast_ref::<NdBroadcastExec>() else {
return Ok(Transformed::no(node));
};
// A `fetch` caps the rows that the filter returns. The nd filter
// records a grid selection and holds no cap. A rewrite that drops
// the `FilterExec` drops the cap too. So keep such a filter here.
if filter.fetch().is_some() {
return Ok(Transformed::no(node));
}

// Split the predicate and route each conjunct: element-wise ones sink
// into the nd filter, the rest stay in a residual filter above.
Expand All @@ -169,12 +185,33 @@ impl PhysicalOptimizerRule for NdFilterPushdown {
return Ok(Transformed::no(node));
}

let nd_filter = Arc::new(NdFilterExec::try_new(broadcast.input().clone(), push)?);
let new_broadcast = Arc::new(NdBroadcastExec::try_new(nd_filter)?);
let nd_filter: Arc<dyn ExecutionPlan> =
Arc::new(NdFilterExec::try_new(broadcast.input().clone(), push)?);

let rewritten: Arc<dyn ExecutionPlan> = if keep.is_empty() {
new_broadcast
// The full predicate sinks, so the projection sinks too. It is
// a plain column list. The nd projection keeps the grid
// selection of the nd filter below it.
let below = match filter.projection().as_deref() {
Some(indices) => Arc::new(NdProjectionExec::try_new_with_schema(
nd_filter,
projected_columns(&filter.input().schema(), indices),
Some(filter.schema()),
)?) as Arc<dyn ExecutionPlan>,
None => nd_filter,
};
Arc::new(NdBroadcastExec::try_new(below)?)
} else {
Arc::new(FilterExec::try_new(conjunction(keep), new_broadcast)?)
// A residual conjunct reads the columns that the projection
// drops. So the projection stays with the residual filter. Build
// the new filter from the original one. This keeps the
// projection, the batch size and the selectivity.
Arc::new(
FilterExecBuilder::from(filter)
.with_predicate(conjunction(keep))
.with_input(Arc::new(NdBroadcastExec::try_new(nd_filter)?))
.build()?,
)
};
Ok(Transformed::yes(rewritten))
})
Expand All @@ -190,6 +227,29 @@ impl PhysicalOptimizerRule for NdFilterPushdown {
}
}

/// Converts the `indices` of a filter projection into `(column, alias)` pairs.
/// `schema` is the input schema of the filter. [`NdProjectionExec`] takes this
/// form.
///
/// A filter projection is always a plain column list. So each output is a
/// [`Column`] that names the field it selects. `FilterExec` validates the indices
/// against the same schema.
fn projected_columns(
schema: &SchemaRef,
indices: &[usize],
) -> Vec<(Arc<dyn PhysicalExpr>, String)> {
indices
.iter()
.map(|&index| {
let name = schema.field(index).name();
(
Arc::new(Column::new(name, index)) as Arc<dyn PhysicalExpr>,
name.clone(),
)
})
.collect()
}

/// Whether an expression can be evaluated before broadcast and give the same
/// result after broadcast — i.e. it is element-wise and deterministic.
///
Expand Down
Loading