fix: apply cheap WHERE conjuncts before inference

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

src/lib.rs

@@ -61,6 +61,7 @@ pub fn register(ctx: &SessionContext, client: Arc<dyn JevClient>) {
     *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();
 }

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<dyn PhysicalExpr>) -> bool {
+    expr.downcast_ref::<ScalarFunctionExpr>()
+        .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<dyn ExecutionPlan>,
+        _: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        plan.transform_up(|plan| {
+            let Some(filter) = plan.downcast_ref::<FilterExec>() 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<Arc<dyn ExecutionPlan>> = vec![];
+            let mut cursor = Arc::clone(filter.input());
+            while cursor.downcast_ref::<RepartitionExec>().is_some() {
+                let child = Arc::clone(cursor.children()[0]);
+                passthrough.push(cursor);
+                cursor = child;
+            }
+            let Some(node) = cursor.downcast_ref::<AsyncFuncExec>() 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<dyn ExecutionPlan> =
+                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<dyn ExecutionPlan>))
+        })
+        .data()
+    }
+}
 #[derive(Debug)]
 pub struct DeduplicateJev;
 impl PhysicalOptimizerRule for DeduplicateJev {

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::<usize>(), 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::<usize>(),
+        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::<usize>(), 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}"
+    );
+}