[DataFusion] Reject negated IN predicates in extractors

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<RangeInclusive<InvocationId>> {
     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<dyn PhysicalExpr>>) -> Arc<dyn PhysicalExpr> {
+        make_in_list(col_name, list, false)
+    }
+
+    fn not_in_list(col_name: &str, list: Vec<Arc<dyn PhysicalExpr>>) -> Arc<dyn PhysicalExpr> {
+        make_in_list(col_name, list, true)
+    }
+
+    fn make_in_list(
+        col_name: &str,
+        list: Vec<Arc<dyn PhysicalExpr>>,
+        negated: bool,
+    ) -> Arc<dyn PhysicalExpr> {
         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<dyn PhysicalExpr>, right: Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
@@ -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"));

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::<Vec<_>>()
+        .join(", ");
+    let records = engine
+        .execute(format!(
+            "SELECT id FROM sys_invocation_status WHERE id NOT IN ({excluded})"
+        ))
+        .await
+        .unwrap()
+        .stream
+        .collect::<Vec<datafusion::common::Result<RecordBatch>>>()
+        .await
+        .remove(0)
+        .unwrap();
+
+    let mut got: Vec<String> = records
+        .column_by_name("id")
+        .unwrap()
+        .as_any()
+        .downcast_ref::<LargeStringArray>()
+        .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();