fix(scan): keep the predicate when a partition set cannot express it

crates/integrations/datafusion/tests/partition_conjunct_pushdown.rs

@@ -0,0 +1,136 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A partition-only filter is pushed down as `Exact`, so DataFusion keeps no
+//! residual. These assert that the partition filter alone enforces the whole
+//! predicate, including conjuncts a partition set cannot express.
+
+mod common;
+
+use common::setup_sql_context;
+use paimon_datafusion::SQLContext;
+
+const TABLE: &str = "paimon.test_db.t_partition_conjuncts";
+
+async fn setup() -> (tempfile::TempDir, SQLContext) {
+    let (tmp, sql_context) = setup_sql_context().await;
+    sql_context
+        .sql(&format!(
+            "CREATE TABLE {TABLE} (dt STRING, id INT) PARTITIONED BY (dt)"
+        ))
+        .await
+        .unwrap();
+    sql_context
+        .sql(&format!(
+            "INSERT INTO {TABLE} VALUES \
+             ('2024-01-01', 1), ('2024-01-02', 2), ('2024-01-03', 3), ('2024-01-04', 4)"
+        ))
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    (tmp, sql_context)
+}
+
+async fn partitions_matching(sql_context: &SQLContext, where_clause: &str) -> Vec<String> {
+    let batches = sql_context
+        .sql(&format!(
+            "SELECT dt FROM {TABLE} WHERE {where_clause} ORDER BY dt"
+        ))
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    let mut out = Vec::new();
+    for batch in &batches {
+        for row in 0..batch.num_rows() {
+            out.push(common::string_value(batch.column(0).as_ref(), row).to_string());
+        }
+    }
+    out
+}
+
+async fn plan_of(sql_context: &SQLContext, where_clause: &str) -> String {
+    let batches = sql_context
+        .sql(&format!(
+            "EXPLAIN SELECT dt FROM {TABLE} WHERE {where_clause}"
+        ))
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    datafusion::arrow::util::pretty::pretty_format_batches(&batches)
+        .unwrap()
+        .to_string()
+}
+
+#[tokio::test]
+async fn contradictory_partition_conjuncts_match_nothing() {
+    let (_tmp, sql_context) = setup().await;
+    let where_clause = "dt = '2024-01-01' AND dt >= '2024-01-02'";
+
+    let plan = plan_of(&sql_context, where_clause).await;
+    assert!(
+        !plan.contains("FilterExec"),
+        "a partition-only filter is Exact both before and after this fix, \
+         so the scan must be the only enforcement:\n{plan}"
+    );
+
+    assert!(
+        partitions_matching(&sql_context, where_clause)
+            .await
+            .is_empty(),
+        "no partition can be both = 2024-01-01 and >= 2024-01-02"
+    );
+}
+
+/// The narrower `=` comes first, so the old second assignment overwrote it with
+/// the wider `IN`. Order matters: with `IN` first the overwrite keeps the `=`
+/// and the answer is accidentally right.
+#[tokio::test]
+async fn an_equality_before_a_wider_in_list_is_not_widened() {
+    let (_tmp, sql_context) = setup().await;
+    let where_clause =
+        "dt = '2024-01-02' AND dt IN ('2024-01-01','2024-01-02','2024-01-03','2024-01-04')";
+
+    assert_eq!(
+        partitions_matching(&sql_context, where_clause).await,
+        vec!["2024-01-02".to_string()]
+    );
+}
+
+#[tokio::test]
+async fn a_range_beside_an_in_list_still_narrows() {
+    let (_tmp, sql_context) = setup().await;
+    let where_clause =
+        "dt >= '2024-01-03' AND dt IN ('2024-01-01','2024-01-02','2024-01-03','2024-01-04')";
+
+    let plan = plan_of(&sql_context, where_clause).await;
+    assert!(
+        !plan.contains("FilterExec"),
+        "a partition-only filter is Exact both before and after this fix, \
+         so the scan must be the only enforcement:\n{plan}"
+    );
+
+    assert_eq!(
+        partitions_matching(&sql_context, where_clause).await,
+        vec!["2024-01-03".to_string(), "2024-01-04".to_string()]
+    );
+}

crates/paimon/src/table/partition_filter.rs

@@ -54,7 +54,9 @@ impl PartitionFilter {
 
         let num_fields = partition_fields.len();
         let mut field_candidates: Vec<Option<Vec<Option<&Datum>>>> = vec![None; num_fields];
-        collect_eq_candidates(&predicate, &mut field_candidates);
+        if !collect_eq_candidates(&predicate, &mut field_candidates) {
+            return PartitionFilter::Predicate(predicate);
+        }
 
         if field_candidates.iter().any(|c| c.is_none()) {
             return PartitionFilter::Predicate(predicate);
@@ -271,36 +273,49 @@ fn build_field_bounds(
     })
 }
 
+/// Collect `Eq`/`In`/`IsNull` candidate values per partition field.
+///
+/// Returns `false` as soon as any node of the tree is not fully represented by
+/// the collected candidates. Callers must then keep the original predicate: a
+/// `PartitionSet` is the sole authority in `matches_entry`, which never looks at
+/// the predicate again, and `ReadBuilder::is_exact_filter_pushdown` lets
+/// DataFusion drop its residual filter for a partition-only predicate.
 fn collect_eq_candidates<'a>(
     predicate: &'a Predicate,
     field_candidates: &mut Vec<Option<Vec<Option<&'a Datum>>>>,
-) {
+) -> bool {
     match predicate {
-        Predicate::And(children) => {
-            for child in children {
-                collect_eq_candidates(child, field_candidates);
-            }
-        }
+        Predicate::And(children) => children
+            .iter()
+            .all(|child| collect_eq_candidates(child, field_candidates)),
         Predicate::Leaf {
             index,
             op,
             literals,
             ..
-        } if *index < field_candidates.len() => match op {
-            PredicateOperator::Eq => {
-                if let Some(lit) = literals.first() {
-                    field_candidates[*index] = Some(vec![Some(lit)]);
-                }
-            }
-            PredicateOperator::In if !literals.is_empty() => {
-                field_candidates[*index] = Some(literals.iter().map(Some).collect());
+        } if *index < field_candidates.len() => {
+            // A second conjunct on the same field used to overwrite the first,
+            // keeping only whichever came last.
+            if field_candidates[*index].is_some() {
+                return false;
             }
-            PredicateOperator::IsNull => {
-                field_candidates[*index] = Some(vec![None]);
+            match op {
+                PredicateOperator::Eq if !literals.is_empty() => {
+                    field_candidates[*index] = Some(vec![Some(&literals[0])]);
+                    true
+                }
+                PredicateOperator::In if !literals.is_empty() => {
+                    field_candidates[*index] = Some(literals.iter().map(Some).collect());
+                    true
+                }
+                PredicateOperator::IsNull => {
+                    field_candidates[*index] = Some(vec![None]);
+                    true
+                }
+                _ => false,
             }
-            _ => {}
-        },
-        _ => {}
+        }
+        _ => false,
     }
 }
 
@@ -460,6 +475,83 @@ mod tests {
         assert!(matches!(filter, PartitionFilter::Predicate(_)));
     }
 
+    fn serialized_dt(fields: &[DataField], dt: &str) -> Vec<u8> {
+        let mut builder = BinaryRowBuilder::new(1);
+        builder.write_datum(0, &Datum::String(dt.into()), fields[0].data_type());
+        builder.build_serialized()
+    }
+
+    /// Coverage is complete, but `>=` is not expressible as a set of values.
+    #[test]
+    fn test_unexpressible_conjunct_on_covered_field_falls_back() {
+        let fields = partition_fields_dt();
+        let pb = PredicateBuilder::new(&fields);
+        let pred = Predicate::and(vec![
+            pb.equal("dt", Datum::String("2024-01-01".into())).unwrap(),
+            pb.greater_or_equal("dt", Datum::String("2024-01-02".into()))
+                .unwrap(),
+        ]);
+        let filter = PartitionFilter::from_predicate(pred, &fields);
+        assert!(matches!(filter, PartitionFilter::Predicate(_)));
+        assert!(!filter
+            .matches_entry(&serialized_dt(&fields, "2024-01-01"))
+            .unwrap());
+    }
+
+    /// Two expressible conjuncts on one field: the second assignment used to
+    /// overwrite the first, keeping whichever came last — here the wider `In`.
+    #[test]
+    fn test_second_conjunct_on_same_field_falls_back() {
+        let fields = partition_fields_dt();
+        let pb = PredicateBuilder::new(&fields);
+        let pred = Predicate::and(vec![
+            pb.equal("dt", Datum::String("2024-01-02".into())).unwrap(),
+            pb.is_in(
+                "dt",
+                vec![
+                    Datum::String("2024-01-01".into()),
+                    Datum::String("2024-01-02".into()),
+                ],
+            )
+            .unwrap(),
+        ]);
+        let filter = PartitionFilter::from_predicate(pred, &fields);
+        assert!(matches!(filter, PartitionFilter::Predicate(_)));
+        assert!(!filter
+            .matches_entry(&serialized_dt(&fields, "2024-01-01"))
+            .unwrap());
+        assert!(filter
+            .matches_entry(&serialized_dt(&fields, "2024-01-02"))
+            .unwrap());
+    }
+
+    /// An `Or` over the partition field narrows the `In` beside it.
+    #[test]
+    fn test_or_conjunct_beside_covering_in_falls_back() {
+        let fields = partition_fields_dt();
+        let pb = PredicateBuilder::new(&fields);
+        let pred = Predicate::and(vec![
+            Predicate::or(vec![
+                pb.equal("dt", Datum::String("2024-01-01".into())).unwrap(),
+                pb.equal("dt", Datum::String("2024-01-02".into())).unwrap(),
+            ]),
+            pb.is_in(
+                "dt",
+                vec![
+                    Datum::String("2024-01-01".into()),
+                    Datum::String("2024-01-02".into()),
+                    Datum::String("2024-01-03".into()),
+                ],
+            )
+            .unwrap(),
+        ]);
+        let filter = PartitionFilter::from_predicate(pred, &fields);
+        assert!(matches!(filter, PartitionFilter::Predicate(_)));
+        assert!(!filter
+            .matches_entry(&serialized_dt(&fields, "2024-01-03"))
+            .unwrap());
+    }
+
     #[test]
     fn test_is_null_in_partition_set() {
         let fields = partition_fields_dt();