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();