diff --git a/README.md b/README.md index d90fac4..434d642 100644 --- a/README.md +++ b/README.md @@ -149,7 +149,10 @@ each row must be judged in complete isolation. `GROUP BY`, `HAVING`, window `OVER (...)` clauses, and join conditions. A call in a join condition must use columns from one side of the join only. - **Filters run first.** Rows removed by `WHERE` are never sent to the service, - including when the call sits inside a CTE or subquery. + including when the call sits inside a CTE or subquery. In a `WHERE` that + combines ordinary conditions with `prompt_jev`, the ordinary conditions are + applied before inference, so `WHERE region = 'eu' AND prompt_jev(...) > 0.8` + only asks about the `eu` rows. - **Limits.** Each row's text may be up to 64 KiB after JSON encoding. Longer text fails the query. Requests are capped at 256 KiB. - **Outages.** Each request is tried three times with a 30-second timeout. If all diff --git a/src/lib.rs b/src/lib.rs index e068aa0..cd2b5f5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -61,6 +61,7 @@ pub fn register(ctx: &SessionContext, client: Arc) { *state = datafusion::execution::SessionStateBuilder::new_from_existing(state.clone()) .with_analyzer_rule(Arc::new(planner::HoistJev)) .with_optimizer_rules(optimizer_rules) + .with_physical_optimizer_rule(Arc::new(optimizer::FilterBeforeJev)) .with_physical_optimizer_rule(Arc::new(optimizer::DeduplicateJev)) .build(); } diff --git a/src/optimizer.rs b/src/optimizer.rs index eb5e25e..b6f2369 100644 --- a/src/optimizer.rs +++ b/src/optimizer.rs @@ -1,17 +1,116 @@ -//! DataFusion 55 extracts every occurrence of an async function, including -//! repeated struct field access. Keep one Jev evaluation and project its result -//! into the original slots so all downstream column indices remain valid. +//! Physical-plan rules around DataFusion 55's `AsyncFuncExec`. +//! +//! * [`DeduplicateJev`]: DataFusion extracts every occurrence of an async +//! function, including repeated struct field access. Keep one Jev evaluation +//! and project its result into the original slots so downstream column +//! indices remain valid. +//! * [`FilterBeforeJev`]: a `WHERE cheap AND prompt_jev(...) > x` is planned as +//! one filter above the async node, so every row is sent for inference before +//! either conjunct runs. Split it so the conjuncts that need no inference run +//! first, beneath the async node. use datafusion::{ common::{ Result, tree_node::{Transformed, TransformedResult, TreeNode}, }, config::ConfigOptions, - physical_expr::{PhysicalExpr, ScalarFunctionExpr, expressions::Column}, + physical_expr::{ + PhysicalExpr, ScalarFunctionExpr, conjunction, expressions::Column, split_conjunction, + utils::collect_columns, + }, physical_optimizer::PhysicalOptimizerRule, - physical_plan::{ExecutionPlan, async_func::AsyncFuncExec, projection::ProjectionExec}, + physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, + async_func::AsyncFuncExec, + filter::{FilterExec, FilterExecBuilder}, + projection::ProjectionExec, + repartition::RepartitionExec, + }, }; use std::sync::Arc; + +fn is_jev_expr(expr: &Arc) -> bool { + expr.downcast_ref::() + .is_some_and(|f| f.fun().name() == "__datafusion_jev") +} + +#[derive(Debug)] +pub struct FilterBeforeJev; +impl PhysicalOptimizerRule for FilterBeforeJev { + fn name(&self) -> &str { + "datafusion_jev_filter_first" + } + fn schema_check(&self) -> bool { + true + } + fn optimize( + &self, + plan: Arc, + _: &ConfigOptions, + ) -> Result> { + plan.transform_up(|plan| { + let Some(filter) = plan.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + // Distribution rules may have placed repartition nodes between the + // filter and the async node; look through them and put them back + // in the same order afterwards. + let mut passthrough: Vec> = vec![]; + let mut cursor = Arc::clone(filter.input()); + while cursor.downcast_ref::().is_some() { + let child = Arc::clone(cursor.children()[0]); + passthrough.push(cursor); + cursor = child; + } + let Some(node) = cursor.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + #[allow(deprecated)] + let async_exprs = node.async_exprs(); + if !async_exprs.iter().any(|e| is_jev_expr(&e.func)) { + return Ok(Transformed::no(plan)); + } + // Columns at or past the async node's input width are async results. + let input_len = node.input().schema().fields().len(); + let (needs_inference, cheap): (Vec<_>, Vec<_>) = split_conjunction(filter.predicate()) + .into_iter() + .cloned() + .partition(|c| { + collect_columns(c) + .iter() + .any(|col| col.index() >= input_len) + }); + if cheap.is_empty() || needs_inference.is_empty() { + return Ok(Transformed::no(plan)); + } + let pre = + FilterExecBuilder::new(conjunction(cheap), Arc::clone(node.input())).build()?; + let mut rebuilt: Arc = + Arc::new(AsyncFuncExec::try_new(async_exprs.to_vec(), Arc::new(pre))?); + for wrapper in passthrough.into_iter().rev() { + rebuilt = wrapper.replace_children( + vec![rebuilt], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + } + // Carry the original filter's fetch (LimitPushdown may have folded a + // LIMIT into it), batch size, and selectivity onto the rebuilt one. + let post = FilterExecBuilder::new(conjunction(needs_inference), rebuilt) + .with_fetch(filter.fetch()) + .with_batch_size(filter.batch_size()) + .with_default_selectivity(filter.default_selectivity()) + .apply_projection( + filter + .projection() + .as_ref() + .map(|p| p.iter().copied().collect()), + )? + .build()?; + Ok(Transformed::yes(Arc::new(post) as Arc)) + }) + .data() + } +} #[derive(Debug)] pub struct DeduplicateJev; impl PhysicalOptimizerRule for DeduplicateJev { diff --git a/tests/sql.rs b/tests/sql.rs index 313117c..b55b2a0 100644 --- a/tests/sql.rs +++ b/tests/sql.rs @@ -778,3 +778,75 @@ async fn two_sided_join_condition_is_a_clear_planning_error() { assert!(!err.contains("called directly"), "{err}"); assert!(mock.calls.lock().unwrap().is_empty()); } + +#[tokio::test] +async fn cheap_conjuncts_filter_rows_before_inference() { + // `WHERE id = 1 AND prompt_jev(...) > 0.5` must not send rows 2 and 3. + let mock = Arc::new(Mock::default()); + let result = run( + mock.clone(), + "SELECT id FROM (VALUES (1,'keep'),(2,'drop'),(3,'drop')) t(id, body) \ + WHERE id = 1 AND prompt_jev(body, 'Urgent?') > 0.5", + ) + .await + .unwrap(); + assert_eq!(result.iter().map(|b| b.num_rows()).sum::(), 1); + let calls = mock.calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!( + calls[0]["questions"].as_object().unwrap().len(), + 1, + "only the row passing the cheap predicate is asked about: {}", + calls[0] + ); +} + +#[tokio::test] +async fn cheap_conjunct_in_a_count_filter_runs_first() { + let mock = Arc::new(Mock::default()); + let result = run( + mock.clone(), + "SELECT count(*) AS n FROM (VALUES ('a','x'),('b','y'),('c','z')) t(name, body) \ + WHERE name <> 'c' AND prompt_jev(body, 'Urgent?') > 0.5", + ) + .await + .unwrap(); + assert!(display(&result).contains('2'), "{}", display(&result)); + let calls = mock.calls.lock().unwrap(); + assert_eq!( + calls + .iter() + .map(|c| c["questions"].as_object().unwrap().len()) + .sum::(), + 2 + ); +} + +#[tokio::test] +async fn a_folded_limit_survives_the_filter_split() { + // With one partition, LimitPushdown folds LIMIT into the FilterExec's fetch. + // The split must keep it, or the query returns every matching row. + let mock = Arc::new(Mock::default()); + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + datafusion_jev::register(&ctx, mock.clone()); + let df = datafusion_jev::sql( + &ctx, + "SELECT id FROM (VALUES (1,'a'),(2,'b'),(3,'c'),(4,'d')) t(id, body) \ + WHERE id > 1 AND prompt_jev(body, 'Urgent?') > 0.5 LIMIT 1", + ) + .await + .unwrap(); + let result = df.collect().await.unwrap(); + assert_eq!(result.iter().map(|b| b.num_rows()).sum::(), 1); + let asked: usize = mock + .calls + .lock() + .unwrap() + .iter() + .map(|c| c["questions"].as_object().unwrap().len()) + .sum(); + assert!( + asked <= 3, + "rows failing `id > 1` must not be asked about: {asked}" + ); +}