From f9b7e871cf07500e9f03f580f1f83a091cce95e9 Mon Sep 17 00:00:00 2001 From: Ahmed Farghal Date: Wed, 15 Jul 2026 16:21:54 +0100 Subject: [PATCH] [DataFusion] Reject negated IN predicates in extractors NOT IN (...) and != predicates enumerate excluded values, but the partition-key extractor and the entry-id/invocation-id filters ignored InList::negated and built lookup sets or ranges from those excluded IDs. Reject negated IN lists in those extractors so DataFusion keeps them as normal filters instead of treating them as positive lookup inputs. Cover the sys_invocation_status regression with excluded and retained invocation IDs on different partition keys, and keep sys_vqueue_entry_status coverage for the same NOT IN behavior. --- crates/storage-query-datafusion/src/filter.rs | 45 +++++++++++++- crates/storage-query-datafusion/src/tests.rs | 62 +++++++++++++++++++ 2 files changed, 104 insertions(+), 3 deletions(-) diff --git a/crates/storage-query-datafusion/src/filter.rs b/crates/storage-query-datafusion/src/filter.rs index e6ae32e49c..596f830f80 100644 --- a/crates/storage-query-datafusion/src/filter.rs +++ b/crates/storage-query-datafusion/src/filter.rs @@ -221,7 +221,9 @@ where continue; }; - if inlist.col.name() != self.column_name { + // A negated list (`NOT IN`/`!=`) enumerates excluded values, so it + // cannot narrow the partition scan. + if inlist.col.name() != self.column_name || inlist.negated { continue; } @@ -481,7 +483,9 @@ fn parse_invocation_id_range( ) -> Option> { let in_list = InList::parse(predicate, 5)?; - if in_list.col.name() != column_name { + // A negated list (`NOT IN`/`!=`) enumerates excluded IDs; using it to build + // a lookup range would fetch exactly the rows that must be filtered out. + if in_list.col.name() != column_name || in_list.negated { return None; } @@ -556,9 +560,21 @@ mod tests { } fn in_list(col_name: &str, list: Vec>) -> Arc { + make_in_list(col_name, list, false) + } + + fn not_in_list(col_name: &str, list: Vec>) -> Arc { + make_in_list(col_name, list, true) + } + + fn make_in_list( + col_name: &str, + list: Vec>, + negated: bool, + ) -> Arc { use datafusion::arrow::datatypes::{DataType, Field, Schema}; let schema = Schema::new(vec![Field::new(col_name, DataType::LargeUtf8, true)]); - Arc::new(InListExpr::try_new(col(col_name), list, false, &schema).expect("valid in-list")) + Arc::new(InListExpr::try_new(col(col_name), list, negated, &schema).expect("valid in-list")) } fn and(left: Arc, right: Arc) -> Arc { @@ -991,6 +1007,29 @@ mod tests { assert!(filter.invocation_ids.is_none()); } + #[test] + fn partition_key_extractor_rejects_negated_in_list() { + let extractor = + FirstMatchingPartitionKeyExtractor::default().with_service_key("service_key"); + + let got = extractor + .try_extract(&[not_in_list("service_key", vec![utf8_lit("key-1")])]) + .expect("extract"); + + assert_eq!(None, got); + } + + #[test] + fn invocation_id_filter_rejects_negated_in_list() { + let id = make_invocation_id("key-1"); + let filter = InvocationIdFilter::new( + FULL_RANGE, + Some(not_in_list("id", vec![utf8_lit(id.to_string())])), + ); + + assert!(filter.invocation_ids.is_none()); + } + #[test] fn vqueue_filter_single_stage_eq() { let predicate = eq(col("stage"), utf8_lit("running")); diff --git a/crates/storage-query-datafusion/src/tests.rs b/crates/storage-query-datafusion/src/tests.rs index 6da1ed19fb..a37626ee93 100644 --- a/crates/storage-query-datafusion/src/tests.rs +++ b/crates/storage-query-datafusion/src/tests.rs @@ -479,6 +479,68 @@ async fn query_sys_invocation_status_completed() { ); } +/// Regression: `NOT IN (> 3 values)` must not be used to build a lookup range +/// from the excluded invocation IDs. Four values is the smallest list that +/// survives DataFusion's inline-list simplifier as a negated `InListExpr`. +#[restate_core::test(flavor = "multi_thread", worker_threads = 2)] +async fn query_sys_invocation_status_not_in() { + use datafusion::arrow::array::Array; + + let invocation_target = InvocationTarget::mock_service(); + let mut engine = MockQueryEngine::create().await; + + let mut ids = Vec::new(); + let mut tx = engine.partition_store().transaction(); + for i in 1..=6u64 { + let id = InvocationId::from_parts(i, InvocationUuid::from_u128(i as u128)); + tx.put_invocation_status( + &id, + &InvocationStatus::Completed(CompletedInvocation { + invocation_target: invocation_target.clone(), + ..CompletedInvocation::mock_neo() + }), + ) + .unwrap(); + ids.push(id.to_string()); + } + tx.commit().await.unwrap(); + drop(tx); + + let excluded = ids[0..4] + .iter() + .map(|id| format!("'{id}'")) + .collect::>() + .join(", "); + let records = engine + .execute(format!( + "SELECT id FROM sys_invocation_status WHERE id NOT IN ({excluded})" + )) + .await + .unwrap() + .stream + .collect::>>() + .await + .remove(0) + .unwrap(); + + let mut got: Vec = records + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(str::to_string) + .collect(); + got.sort(); + + let mut expected = vec![ids[4].clone(), ids[5].clone()]; + expected.sort(); + + assert_eq!(got, expected); +} + #[restate_core::test(flavor = "multi_thread", worker_threads = 2)] async fn query_sys_invocation_suspended_waiting() { let invocation_id = InvocationId::mock_random();