From 6ff9a077b458ec09db42fee553ab5a4270a49308 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 18 Aug 2026 10:26:05 +0200 Subject: [PATCH] Sink a filter that holds a projection into NdFilterExec (#397) NdFilterPushdown refuses a FilterExec that holds a projection. DataFusion adds a projection when the select list is narrower than the predicate. count(*) gets an empty projection. So the most common spatial query filters after the broadcast. Split the filter. The conjuncts go into NdFilterExec. The projection becomes an NdProjectionExec above it, below the broadcast. A residual conjunct reads the columns that the projection drops, so the projection stays with the residual filter above the broadcast. Keep a filter that holds a fetch. The nd filter has no cap. The old rule dropped the whole node, and the cap with it. Add four unit tests in nd/mod.rs. Each test compares the rewritten plan against the rows of the original plan. Assisted-by: Claude:claude-opus-5 --- beacon-db/beacon-datafusion-ext/src/nd/mod.rs | 204 ++++++++++++++++++ .../beacon-datafusion-ext/src/nd/optimizer.rs | 86 ++++++-- 2 files changed, 277 insertions(+), 13 deletions(-) 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. ///