diff --git a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs index 7cee86e4..49bd5519 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs @@ -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 = + binary(col("lon", &schema).unwrap(), Operator::Eq, lit(5i32), &schema).unwrap(); + + let original: Arc = 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 = + binary(col("lon", &schema).unwrap(), Operator::Eq, lit(5i32), &schema).unwrap(); + + let original: Arc = 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 = 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 = 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 = + binary(col("lon", &schema).unwrap(), Operator::Eq, lit(5i32), &schema).unwrap(); + + let original: Arc = 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); + } } diff --git a/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs b/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs index 2b16818e..bfbb24a1 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs @@ -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; @@ -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; @@ -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; @@ -144,15 +160,15 @@ impl PhysicalOptimizerRule for NdFilterPushdown { let Some(filter) = node.as_any().downcast_ref::() 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::() 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. @@ -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 = + Arc::new(NdFilterExec::try_new(broadcast.input().clone(), push)?); + let rewritten: Arc = 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, + 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)) }) @@ -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, String)> { + indices + .iter() + .map(|&index| { + let name = schema.field(index).name(); + ( + Arc::new(Column::new(name, index)) as Arc, + 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. ///