query/dist_plan/
predicate_extractor.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Predicate extraction for partition pruning.
16//!
17//! [`PredicateExtractor`] extract a list of [`PartitionExpr`] from given [`LogicalPlan`].
18
19use std::collections::HashSet;
20
21use common_telemetry::debug;
22use datafusion_common::Result as DfResult;
23use datafusion_expr::{Expr, LogicalPlan, Operator, utils::split_conjunction};
24use datatypes::value::Value;
25use partition::expr::{Operand, PartitionExpr, RestrictedOp};
26
27/// Extracts a list of [`PartitionExpr`] from given [`LogicalPlan`]
28pub struct PredicateExtractor;
29
30impl PredicateExtractor {
31    /// Extract partition expressions for partition columns from logical plan  
32    /// This method returns PartitionExpr objects suitable for ConstraintPruner
33    pub fn extract_partition_expressions(
34        plan: &LogicalPlan,
35        partition_columns: &[String],
36    ) -> DfResult<Vec<PartitionExpr>> {
37        // Collect all filter expressions from the logical plan
38        let mut filter_exprs = Vec::new();
39        Self::collect_filter_expressions(plan, &mut filter_exprs)?;
40
41        if filter_exprs.is_empty() {
42            return Ok(Vec::new());
43        }
44
45        // Convert each DataFusion filter expression to PartitionExpr
46        let mut partition_exprs = Vec::with_capacity(filter_exprs.len());
47        let partition_set: HashSet<String> = partition_columns.iter().cloned().collect();
48
49        for filter_expr in filter_exprs {
50            let partition_expr = match DataFusionExprConverter::convert(&filter_expr) {
51                Ok(expr) => Some(expr),
52                Err(err) => {
53                    debug!(
54                        "Failed to convert filter expression to PartitionExpr: {}, trying conjuncts",
55                        err
56                    );
57                    split_conjunction(&filter_expr)
58                        .into_iter()
59                        .filter_map(|conjunct| match DataFusionExprConverter::convert(conjunct) {
60                            Ok(expr) => Some(expr),
61                            Err(err) => {
62                                debug!(
63                                    "Failed to convert filter conjunct to PartitionExpr: {}, skipping",
64                                    err
65                                );
66                                None
67                            }
68                        })
69                        // Keep salvaged conjuncts in one AND expression, not as alternatives.
70                        .reduce(PartitionExpr::and)
71                }
72            };
73            let Some(partition_expr) = partition_expr else {
74                continue;
75            };
76
77            // Check expression for safe partition pruning
78            match ExpressionChecker::check_expression_for_pruning(&partition_expr, &partition_set) {
79                ExpressionCheckResult::UseAsIs(expr) => {
80                    partition_exprs.push(expr);
81                }
82                ExpressionCheckResult::UsePartial(exprs) => {
83                    partition_exprs.extend(exprs);
84                }
85                ExpressionCheckResult::Drop => {
86                    debug!(
87                        "Dropping mixed expression for correctness: {}",
88                        partition_expr
89                    );
90                }
91            }
92        }
93
94        debug!(
95            "Extracted {} partition expressions from logical plan for partition columns: {:?}",
96            partition_exprs.len(),
97            partition_columns
98        );
99
100        Ok(partition_exprs)
101    }
102
103    /// Collect all filter expressions from a logical plan.
104    ///
105    /// Besides explicit [`LogicalPlan::Filter`] nodes, this must also collect
106    /// predicates already stored in [`LogicalPlan::TableScan`] filters. The
107    /// distributed planner runs a focused DataFusion `PushDownFilter` pass
108    /// before `MergeScan` wrapping, so partition predicates may no longer exist
109    /// as standalone `Filter` nodes by the time region pruning calls this
110    /// extractor. If we ignored `TableScan.filters`, region pruning would miss
111    /// predicates that were successfully pushed down for scan-level pruning.
112    fn collect_filter_expressions(plan: &LogicalPlan, expressions: &mut Vec<Expr>) -> DfResult<()> {
113        if let LogicalPlan::Filter(filter) = plan {
114            expressions.push(filter.predicate.clone());
115        }
116
117        // Collect filters that DataFusion's PushDownFilter stored in TableScan.
118        // `TableScan.filters` is conjunctive: DataFusion passes scan filters as
119        // a list but the table scan must satisfy all of them. Preserve that AND
120        // semantics for partition pruning instead of returning the filters as
121        // independent top-level expressions.
122        if let LogicalPlan::TableScan(table_scan) = plan
123            && let Some(expr) = Self::conjunction(table_scan.filters.iter().cloned())
124        {
125            expressions.push(expr);
126        }
127
128        // Recursively visit children
129        for child in plan.inputs() {
130            Self::collect_filter_expressions(child, expressions)?;
131        }
132
133        // TODO(ruihang): support plans involve multiple relations.
134        if plan.inputs().len() > 1 {
135            expressions.clear();
136        }
137
138        Ok(())
139    }
140
141    fn conjunction(mut expressions: impl Iterator<Item = Expr>) -> Option<Expr> {
142        let first = expressions.next()?;
143        Some(expressions.fold(first, |acc, expr| acc.and(expr)))
144    }
145}
146
147/// Result of analyzing an expression for partition pruning safety
148#[derive(Debug, Clone)]
149enum ExpressionCheckResult {
150    /// Expression is safe to use as-is for pruning (only involves partition columns)
151    UseAsIs(PartitionExpr),
152    /// Extract only these parts for AND expressions (mixed with non-partition columns)
153    UsePartial(Vec<PartitionExpr>),
154    /// Drop the entire expression (unsafe for pruning, e.g., OR with non-partition columns)
155    Drop,
156}
157
158/// Checks expressions to determine safe pruning strategies for mixed partition/non-partition expressions
159struct ExpressionChecker;
160
161impl ExpressionChecker {
162    /// Check a partition expression to determine how it should be handled for safe pruning
163    fn check_expression_for_pruning(
164        expr: &PartitionExpr,
165        partition_columns: &HashSet<String>,
166    ) -> ExpressionCheckResult {
167        match expr.op() {
168            RestrictedOp::And => {
169                // For AND expressions, we can extract only the partition-related parts
170                // because if any part fails, the entire AND fails
171                let mut partition_constraints = Vec::new();
172                Self::extract_and_constraints(expr, partition_columns, &mut partition_constraints);
173
174                if partition_constraints.is_empty() {
175                    ExpressionCheckResult::Drop
176                } else if Self::expr_only_involves_partition_columns(expr, partition_columns) {
177                    // If the entire expression only involves partition columns, use as-is
178                    ExpressionCheckResult::UseAsIs(expr.clone())
179                } else {
180                    // Mixed expression: return only the partition parts
181                    ExpressionCheckResult::UsePartial(partition_constraints)
182                }
183            }
184            RestrictedOp::Or => {
185                // For OR expressions, we can only use them if ALL branches involve only partition columns
186                // because if any branch involves non-partition columns, we cannot safely prune
187                if Self::expr_only_involves_partition_columns(expr, partition_columns) {
188                    ExpressionCheckResult::UseAsIs(expr.clone())
189                } else {
190                    // Mixed OR expression: must drop entirely for safety
191                    ExpressionCheckResult::Drop
192                }
193            }
194            _ => {
195                // For comparison operations (=, <, >, etc.), check if they only involve partition columns
196                if Self::expr_only_involves_partition_columns(expr, partition_columns) {
197                    ExpressionCheckResult::UseAsIs(expr.clone())
198                } else {
199                    ExpressionCheckResult::Drop
200                }
201            }
202        }
203    }
204
205    /// Extract partition-related constraints from AND expressions recursively
206    fn extract_and_constraints(
207        expr: &PartitionExpr,
208        partition_columns: &HashSet<String>,
209        result: &mut Vec<PartitionExpr>,
210    ) {
211        if let RestrictedOp::And = expr.op() {
212            // Recursively process both sides of AND
213            Self::extract_constraints_from_operand(expr.lhs(), partition_columns, result);
214            Self::extract_constraints_from_operand(expr.rhs(), partition_columns, result);
215        } else {
216            // Non-AND expression: check if it involves partition columns
217            if Self::expr_only_involves_partition_columns(expr, partition_columns) {
218                result.push(expr.clone());
219            }
220        }
221    }
222
223    /// Extract constraints from an operand (which might be a column, value, or nested expression)
224    fn extract_constraints_from_operand(
225        operand: &Operand,
226        partition_columns: &HashSet<String>,
227        result: &mut Vec<PartitionExpr>,
228    ) {
229        match operand {
230            Operand::Column(_) | Operand::Value(_) | Operand::IcebergHash(_) => {
231                // This shouldn't happen in well-formed expressions
232            }
233            Operand::Expr(expr) => {
234                Self::extract_and_constraints(expr, partition_columns, result);
235            }
236        }
237    }
238
239    /// Check if an expression involves ONLY partition columns
240    fn expr_only_involves_partition_columns(
241        expr: &PartitionExpr,
242        partition_columns: &HashSet<String>,
243    ) -> bool {
244        Self::operand_only_involves_partition_columns(expr.lhs(), partition_columns)
245            && Self::operand_only_involves_partition_columns(expr.rhs(), partition_columns)
246    }
247
248    /// Check if an operand involves ONLY partition columns or values
249    fn operand_only_involves_partition_columns(
250        operand: &Operand,
251        partition_columns: &HashSet<String>,
252    ) -> bool {
253        match operand {
254            Operand::Column(col) => partition_columns.contains(col),
255            Operand::IcebergHash(hash) => partition_columns.contains(hash.column()),
256            Operand::Value(_) => true, // Values are always safe
257            Operand::Expr(expr) => {
258                Self::expr_only_involves_partition_columns(expr, partition_columns)
259            }
260        }
261    }
262}
263
264/// Converts DataFusion expressions to PartitionExpr
265struct DataFusionExprConverter;
266
267impl DataFusionExprConverter {
268    /// Convert DataFusion Expr to PartitionExpr
269    pub fn convert(expr: &Expr) -> DfResult<PartitionExpr> {
270        match expr {
271            Expr::BinaryExpr(binary_expr) => {
272                let lhs = Self::convert_to_operand(&binary_expr.left)?;
273                let rhs = Self::convert_to_operand(&binary_expr.right)?;
274                let op = Self::convert_operator(&binary_expr.op)?;
275
276                Ok(PartitionExpr::new(lhs, op, rhs))
277            }
278            Expr::InList(inlist_expr) => {
279                // Convert col IN (val1, val2, val3) to col = val1 OR col = val2 OR col = val3
280                // Handle negation: col NOT IN (val1, val2) to col != val1 AND col != val2
281                let column_operand = Self::convert_to_operand(&inlist_expr.expr)?;
282
283                if inlist_expr.list.is_empty() {
284                    return Err(datafusion_common::DataFusionError::Plan(
285                        "InList with empty list is not supported".to_string(),
286                    ));
287                }
288
289                let op = if inlist_expr.negated {
290                    RestrictedOp::NotEq
291                } else {
292                    RestrictedOp::Eq
293                };
294
295                let connector_op = if inlist_expr.negated {
296                    RestrictedOp::And // NOT IN becomes col != val1 AND col != val2
297                } else {
298                    RestrictedOp::Or // IN becomes col = val1 OR col = val2
299                };
300
301                // Convert each value in the list to an equality/inequality expression
302                let mut expressions = Vec::new();
303                for value_expr in &inlist_expr.list {
304                    let value_operand = Self::convert_to_operand(value_expr)?;
305                    expressions.push(PartitionExpr::new(
306                        column_operand.clone(),
307                        op.clone(),
308                        value_operand,
309                    ));
310                }
311
312                // Chain expressions with OR/AND
313                let mut expr_iter = expressions.into_iter();
314                let mut result = expr_iter.next().unwrap();
315                for expr in expr_iter {
316                    result = PartitionExpr::new(
317                        Operand::Expr(result),
318                        connector_op.clone(),
319                        Operand::Expr(expr),
320                    );
321                }
322
323                Ok(result)
324            }
325            Expr::Between(between_expr) => {
326                // Convert col BETWEEN low AND high to col >= low AND col <= high
327                // Handle negation: col NOT BETWEEN low AND high to col < low OR col > high
328                let column_operand = Self::convert_to_operand(&between_expr.expr)?;
329                let low_operand = Self::convert_to_operand(&between_expr.low)?;
330                let high_operand = Self::convert_to_operand(&between_expr.high)?;
331
332                if between_expr.negated {
333                    // NOT BETWEEN: col < low OR col > high
334                    let left_expr =
335                        PartitionExpr::new(column_operand.clone(), RestrictedOp::Lt, low_operand);
336                    let right_expr =
337                        PartitionExpr::new(column_operand, RestrictedOp::Gt, high_operand);
338                    Ok(PartitionExpr::new(
339                        Operand::Expr(left_expr),
340                        RestrictedOp::Or,
341                        Operand::Expr(right_expr),
342                    ))
343                } else {
344                    // BETWEEN: col >= low AND col <= high
345                    let left_expr =
346                        PartitionExpr::new(column_operand.clone(), RestrictedOp::GtEq, low_operand);
347                    let right_expr =
348                        PartitionExpr::new(column_operand, RestrictedOp::LtEq, high_operand);
349                    Ok(PartitionExpr::new(
350                        Operand::Expr(left_expr),
351                        RestrictedOp::And,
352                        Operand::Expr(right_expr),
353                    ))
354                }
355            }
356            Expr::IsNull(expr) => {
357                // Convert col IS NULL to a PartitionExpr
358                let column_operand = Self::convert_to_operand(expr)?;
359                Ok(PartitionExpr::new(
360                    column_operand,
361                    RestrictedOp::Eq,
362                    Operand::Value(Value::Null),
363                ))
364            }
365            Expr::IsNotNull(expr) => {
366                // Convert col IS NOT NULL to a PartitionExpr
367                let column_operand = Self::convert_to_operand(expr)?;
368                Ok(PartitionExpr::new(
369                    column_operand,
370                    RestrictedOp::NotEq,
371                    Operand::Value(Value::Null),
372                ))
373            }
374            Expr::Not(expr) => {
375                // Handle NOT expressions by inverting the inner expression
376                match expr.as_ref() {
377                    Expr::BinaryExpr(binary_expr) => {
378                        let lhs = Self::convert_to_operand(&binary_expr.left)?;
379                        let rhs = Self::convert_to_operand(&binary_expr.right)?;
380                        let inverted_op = Self::invert_operator(&binary_expr.op)?;
381
382                        Ok(PartitionExpr::new(lhs, inverted_op, rhs))
383                    }
384                    Expr::IsNull(inner_expr) => {
385                        // NOT (col IS NULL) becomes col IS NOT NULL
386                        let column_operand = Self::convert_to_operand(inner_expr)?;
387                        Ok(PartitionExpr::new(
388                            column_operand,
389                            RestrictedOp::NotEq,
390                            Operand::Value(Value::Null),
391                        ))
392                    }
393                    Expr::IsNotNull(inner_expr) => {
394                        // NOT (col IS NOT NULL) becomes col IS NULL
395                        let column_operand = Self::convert_to_operand(inner_expr)?;
396                        Ok(PartitionExpr::new(
397                            column_operand,
398                            RestrictedOp::Eq,
399                            Operand::Value(Value::Null),
400                        ))
401                    }
402                    _ => {
403                        debug!(
404                            "Unsupported NOT expression for partition pruning: {:?}",
405                            expr
406                        );
407                        Err(datafusion_common::DataFusionError::Plan(format!(
408                            "NOT expression with inner type {:?} not supported for partition pruning",
409                            expr
410                        )))
411                    }
412                }
413            }
414            _ => Err(datafusion_common::DataFusionError::Plan(format!(
415                "Unsupported expression type for conversion: {:?}",
416                expr
417            ))),
418        }
419    }
420
421    /// Convert DataFusion Expr to Operand
422    fn convert_to_operand(expr: &Expr) -> DfResult<Operand> {
423        match expr {
424            Expr::Column(col) => {
425                // Handle qualified column names (table.column) by extracting just the column name
426                // For partition pruning, we typically only care about the column name itself
427                let column_name = if let Some(relation) = &col.relation {
428                    debug!(
429                        "Using qualified column reference: {}.{}",
430                        relation, col.name
431                    );
432                    col.name.clone()
433                } else {
434                    col.name.clone()
435                };
436                Ok(Operand::Column(column_name))
437            }
438            Expr::Literal(scalar_value, _) => {
439                let value = Value::try_from(scalar_value.clone()).map_err(|err| {
440                    datafusion_common::DataFusionError::Plan(format!(
441                        "Failed to convert literal for partition pruning: {err}"
442                    ))
443                })?;
444                Ok(Operand::Value(value))
445            }
446            Expr::Alias(alias_expr) => {
447                // Unwrap alias to get the actual expression
448                Self::convert_to_operand(&alias_expr.expr)
449            }
450            Expr::Cast(cast_expr) => {
451                // Column casts need input type/value equivalence proof; literals must use the actual cast value.
452                if let Expr::Literal(value, _) = cast_expr.expr.as_ref() {
453                    let value = value.cast_to(cast_expr.field.data_type())?;
454                    let value = Value::try_from(value).map_err(|err| {
455                        datafusion_common::DataFusionError::Plan(format!(
456                            "Failed to convert cast literal for partition pruning: {err}"
457                        ))
458                    })?;
459                    Ok(Operand::Value(value))
460                } else {
461                    debug!("Skipping cast expression for partition pruning: {:?}", expr);
462                    Err(datafusion_common::DataFusionError::Plan(
463                        "Cast expression not supported for partition pruning".to_string(),
464                    ))
465                }
466            }
467            other => {
468                let partition_expr = Self::convert(other)?;
469                Ok(Operand::Expr(partition_expr))
470            }
471        }
472    }
473
474    /// Convert DataFusion Operator to RestrictedOp
475    fn convert_operator(op: &Operator) -> DfResult<RestrictedOp> {
476        match op {
477            Operator::Eq => Ok(RestrictedOp::Eq),
478            Operator::NotEq => Ok(RestrictedOp::NotEq),
479            Operator::Lt => Ok(RestrictedOp::Lt),
480            Operator::LtEq => Ok(RestrictedOp::LtEq),
481            Operator::Gt => Ok(RestrictedOp::Gt),
482            Operator::GtEq => Ok(RestrictedOp::GtEq),
483            Operator::And => Ok(RestrictedOp::And),
484            Operator::Or => Ok(RestrictedOp::Or),
485            _ => Err(datafusion_common::DataFusionError::Plan(format!(
486                "Unsupported operator: {:?}",
487                op
488            ))),
489        }
490    }
491
492    /// Invert a DataFusion Operator for NOT expressions
493    fn invert_operator(op: &Operator) -> DfResult<RestrictedOp> {
494        let Some(negated) = op.negate() else {
495            return Err(datafusion_common::DataFusionError::Plan(format!(
496                "Cannot invert operator: {:?}",
497                op
498            )));
499        };
500        Self::convert_operator(&negated)
501    }
502}
503
504#[cfg(test)]
505mod tests {
506    use std::sync::Arc;
507
508    use datafusion::arrow::datatypes::{DataType, Field, Schema};
509    use datafusion::common::Column;
510    use datafusion::datasource::DefaultTableSource;
511    use datafusion_expr::{LogicalPlanBuilder, col, lit};
512    use datatypes::value::Value;
513    use partition::expr::{Operand, PartitionExpr, RestrictedOp};
514
515    use super::*;
516
517    fn create_test_table_scan() -> LogicalPlan {
518        let schema = Arc::new(Schema::new(vec![
519            Field::new(
520                "timestamp",
521                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
522                false,
523            ),
524            Field::new("user_id", DataType::Int64, false),
525            Field::new("value", DataType::Int64, false),
526        ]));
527
528        let empty_table = datafusion::datasource::empty::EmptyTable::new(schema);
529        let table_source = Arc::new(DefaultTableSource::new(Arc::new(empty_table)));
530
531        LogicalPlanBuilder::scan("test", table_source, None)
532            .unwrap()
533            .build()
534            .unwrap()
535    }
536
537    struct FilterTestCase {
538        name: &'static str,
539        filter_expr: Expr,
540        expected_partition_exprs: Vec<PartitionExpr>,
541        partition_columns: Vec<&'static str>,
542    }
543
544    impl FilterTestCase {
545        fn new(
546            name: &'static str,
547            filter_expr: Expr,
548            expected_partition_exprs: Vec<PartitionExpr>,
549            partition_columns: Vec<&'static str>,
550        ) -> Self {
551            Self {
552                name,
553                filter_expr,
554                expected_partition_exprs,
555                partition_columns,
556            }
557        }
558    }
559
560    /// Helper to check partition expressions for a set of test cases.
561    fn check_partition_expressions(cases: Vec<FilterTestCase>) {
562        for case in cases {
563            let table_scan = create_test_table_scan();
564            let filter = case.filter_expr.clone();
565
566            let plan = LogicalPlanBuilder::from(table_scan)
567                .filter(filter)
568                .unwrap()
569                .build()
570                .unwrap();
571
572            let partition_columns: Vec<String> = case
573                .partition_columns
574                .iter()
575                .map(|s| s.to_string())
576                .collect();
577            let partition_exprs =
578                PredicateExtractor::extract_partition_expressions(&plan, &partition_columns)
579                    .unwrap();
580            let expected = case.expected_partition_exprs.clone();
581            assert_eq!(
582                partition_exprs, expected,
583                "Test case '{}': expected partition expressions {:?}, got {:?}",
584                case.name, expected, partition_exprs
585            );
586        }
587    }
588
589    #[test]
590    fn test_extracts_table_scan_filters() {
591        let table_scan = create_test_table_scan();
592        let filter = col("user_id").gt_eq(lit(100i64));
593        let LogicalPlan::TableScan(scan) = table_scan else {
594            panic!("expected test table scan");
595        };
596        let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan {
597            filters: vec![filter],
598            ..scan
599        });
600
601        let partition_exprs =
602            PredicateExtractor::extract_partition_expressions(&plan, &["user_id".to_string()])
603                .unwrap();
604
605        assert_eq!(
606            partition_exprs,
607            vec![PartitionExpr::new(
608                Operand::Column("user_id".to_string()),
609                RestrictedOp::GtEq,
610                Operand::Value(Value::Int64(100)),
611            )]
612        );
613    }
614
615    #[test]
616    fn test_dictionary_cast_preserves_partition_constraint() {
617        let dictionary_type =
618            DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8));
619        let filter = col("tag").eq(Expr::Cast(datafusion_expr::Cast::new(
620            Box::new(lit("b")),
621            dictionary_type,
622        )));
623
624        let partition_expr = DataFusionExprConverter::convert(&filter).unwrap();
625        assert_eq!(
626            partition_expr,
627            PartitionExpr::new(
628                Operand::Column("tag".to_string()),
629                RestrictedOp::Eq,
630                Operand::Value(Value::String("b".into())),
631            )
632        );
633    }
634
635    #[test]
636    fn test_cast_literals_are_evaluated_and_column_casts_are_skipped() {
637        let casts = [
638            (DataType::Int64, lit(1i64)),
639            (DataType::Float32, lit(1.0f32)),
640            (DataType::Utf8, lit("1")),
641            (DataType::Date32, lit(1i32)),
642            (
643                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Nanosecond, None),
644                lit(1i64),
645            ),
646        ];
647        for (data_type, value) in casts {
648            let cast = Expr::Cast(datafusion_expr::Cast::new(
649                Box::new(col("user_id")),
650                data_type,
651            ));
652            assert!(DataFusionExprConverter::convert(&cast.eq(value)).is_err());
653        }
654
655        let chain = Expr::Cast(datafusion_expr::Cast::new(
656            Box::new(Expr::Cast(datafusion_expr::Cast::new(
657                Box::new(col("user_id")),
658                DataType::Float64,
659            ))),
660            DataType::Int64,
661        ));
662        assert!(DataFusionExprConverter::convert(&chain.eq(lit(1i64))).is_err());
663
664        let cast_int = Expr::Cast(datafusion_expr::Cast::new(
665            Box::new(lit(1.5f64)),
666            DataType::Int64,
667        ));
668        assert_eq!(
669            DataFusionExprConverter::convert(&col("user_id").eq(cast_int))
670                .unwrap()
671                .rhs(),
672            &Operand::Value(Value::Int64(1))
673        );
674
675        let cast_string = Expr::Cast(datafusion_expr::Cast::new(
676            Box::new(lit("12")),
677            DataType::Int64,
678        ));
679        assert_eq!(
680            DataFusionExprConverter::convert(&col("user_id").eq(cast_string))
681                .unwrap()
682                .rhs(),
683            &Operand::Value(Value::Int64(12))
684        );
685
686        let unsupported = Expr::Cast(datafusion_expr::Cast::new(
687            Box::new(lit(1i64)),
688            DataType::Decimal256(10, 2),
689        ));
690        assert!(DataFusionExprConverter::convert_to_operand(&unsupported).is_err());
691    }
692
693    #[test]
694    fn test_cast_predicates_are_skipped_in_and_and_or() {
695        let cast = Expr::Cast(datafusion_expr::Cast::new(
696            Box::new(col("user_id")),
697            DataType::Float32,
698        ));
699        let lower = PartitionExpr::new(
700            Operand::Column("user_id".to_string()),
701            RestrictedOp::GtEq,
702            Operand::Value(Value::Int64(1)),
703        );
704        let upper = PartitionExpr::new(
705            Operand::Column("user_id".to_string()),
706            RestrictedOp::Lt,
707            Operand::Value(Value::Int64(10)),
708        );
709        let bounds = PartitionExpr::new(
710            Operand::Expr(lower.clone()),
711            RestrictedOp::And,
712            Operand::Expr(upper.clone()),
713        );
714        check_partition_expressions(vec![
715            FilterTestCase::new(
716                "cast_in_and_preserves_partition_bound",
717                col("user_id")
718                    .eq(lit(1i64))
719                    .and(cast.clone().eq(lit(1.0f32))),
720                vec![PartitionExpr::new(
721                    Operand::Column("user_id".to_string()),
722                    RestrictedOp::Eq,
723                    Operand::Value(Value::Int64(1)),
724                )],
725                vec!["user_id"],
726            ),
727            FilterTestCase::new(
728                "cast_between_two_bounds_preserves_both",
729                col("user_id")
730                    .gt_eq(lit(1i64))
731                    .and(col("user_id").lt(lit(10i64)))
732                    .and(cast.clone().eq(lit(1.0f32))),
733                vec![bounds.clone()],
734                vec!["user_id"],
735            ),
736            FilterTestCase::new(
737                "cast_left_of_bounds_preserves_both",
738                cast.clone()
739                    .eq(lit(1.0f32))
740                    .and(col("user_id").gt_eq(lit(1i64)))
741                    .and(col("user_id").lt(lit(10i64))),
742                vec![bounds.clone()],
743                vec!["user_id"],
744            ),
745            FilterTestCase::new(
746                "cast_between_bounds_preserves_both",
747                col("user_id")
748                    .gt_eq(lit(1i64))
749                    .and(cast.clone().eq(lit(1.0f32)))
750                    .and(col("user_id").lt(lit(10i64))),
751                vec![bounds],
752                vec!["user_id"],
753            ),
754            FilterTestCase::new(
755                "is_not_null_cast_preserves_partition_bound",
756                col("user_id").eq(lit(1i64)).and(
757                    Expr::Cast(datafusion_expr::Cast::new(
758                        Box::new(col("value")),
759                        DataType::Int32,
760                    ))
761                    .is_not_null(),
762                ),
763                vec![PartitionExpr::new(
764                    Operand::Column("user_id".to_string()),
765                    RestrictedOp::Eq,
766                    Operand::Value(Value::Int64(1)),
767                )],
768                vec!["user_id"],
769            ),
770            FilterTestCase::new(
771                "all_unsupported_conjunction_is_dropped",
772                Expr::Cast(datafusion_expr::Cast::new(
773                    Box::new(col("value")),
774                    DataType::Int32,
775                ))
776                .eq(lit(1i32))
777                .and(
778                    Expr::Cast(datafusion_expr::Cast::new(
779                        Box::new(col("timestamp")),
780                        DataType::Float32,
781                    ))
782                    .eq(lit(1.0f32)),
783                ),
784                vec![],
785                vec!["user_id"],
786            ),
787            FilterTestCase::new(
788                "nested_and_with_or_side_preserves_bound",
789                col("user_id").eq(lit(1i64)).and(
790                    col("user_id")
791                        .eq(lit(2i64))
792                        .or(cast.clone().eq(lit(1.0f32))),
793                ),
794                vec![PartitionExpr::new(
795                    Operand::Column("user_id".to_string()),
796                    RestrictedOp::Eq,
797                    Operand::Value(Value::Int64(1)),
798                )],
799                vec!["user_id"],
800            ),
801            FilterTestCase::new(
802                "or_with_nested_and_cast_drops_whole_expression",
803                col("user_id")
804                    .eq(lit(1i64))
805                    .and(cast.clone().eq(lit(1.0f32)))
806                    .or(col("user_id").eq(lit(2i64))),
807                vec![],
808                vec!["user_id"],
809            ),
810            FilterTestCase::new(
811                "cast_in_or_drops_entire_predicate",
812                col("user_id").eq(lit(1i64)).or(cast.eq(lit(1.0f32))),
813                vec![],
814                vec!["user_id"],
815            ),
816        ]);
817    }
818
819    #[test]
820    fn test_combines_table_scan_filters_as_conjunction() {
821        let table_scan = create_test_table_scan();
822        let filter_a = col("user_id").eq(lit(10i64));
823        let filter_b = col("value").eq(lit(20i64));
824        let LogicalPlan::TableScan(scan) = table_scan else {
825            panic!("expected test table scan");
826        };
827        let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan {
828            filters: vec![filter_a.clone(), filter_b.clone()],
829            ..scan.clone()
830        });
831        let partition_columns = ["user_id".to_string(), "value".to_string()];
832        let expected = vec![PartitionExpr::new(
833            Operand::Expr(PartitionExpr::new(
834                Operand::Column("user_id".to_string()),
835                RestrictedOp::Eq,
836                Operand::Value(Value::Int64(10)),
837            )),
838            RestrictedOp::And,
839            Operand::Expr(PartitionExpr::new(
840                Operand::Column("value".to_string()),
841                RestrictedOp::Eq,
842                Operand::Value(Value::Int64(20)),
843            )),
844        )];
845
846        let partition_exprs =
847            PredicateExtractor::extract_partition_expressions(&plan, &partition_columns).unwrap();
848        assert_eq!(partition_exprs, expected);
849
850        let unsupported = Expr::Cast(datafusion_expr::Cast::new(
851            Box::new(col("timestamp")),
852            DataType::Float32,
853        ))
854        .eq(lit(1.0f32));
855        let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan {
856            filters: vec![filter_a, unsupported, filter_b],
857            ..scan
858        });
859
860        let partition_exprs =
861            PredicateExtractor::extract_partition_expressions(&plan, &partition_columns).unwrap();
862
863        assert_eq!(partition_exprs, expected);
864    }
865
866    #[test]
867    fn test_basic_constraints_extraction() {
868        let cases = vec![
869            FilterTestCase::new(
870                "non_partition_column_ignored",
871                col("value").gt_eq(lit(100i64)),
872                vec![],
873                vec!["user_id"],
874            ),
875            FilterTestCase::new(
876                "simple_constraint",
877                col("user_id").gt_eq(lit(100i64)),
878                vec![PartitionExpr::new(
879                    Operand::Column("user_id".to_string()),
880                    RestrictedOp::GtEq,
881                    Operand::Value(Value::Int64(100)),
882                )],
883                vec!["user_id"],
884            ),
885            FilterTestCase::new(
886                "or_expression",
887                col("user_id")
888                    .eq(lit(100i64))
889                    .or(col("user_id").eq(lit(200i64))),
890                vec![PartitionExpr::new(
891                    Operand::Expr(PartitionExpr::new(
892                        Operand::Column("user_id".to_string()),
893                        RestrictedOp::Eq,
894                        Operand::Value(Value::Int64(100)),
895                    )),
896                    RestrictedOp::Or,
897                    Operand::Expr(PartitionExpr::new(
898                        Operand::Column("user_id".to_string()),
899                        RestrictedOp::Eq,
900                        Operand::Value(Value::Int64(200)),
901                    )),
902                )],
903                vec!["user_id"],
904            ),
905            FilterTestCase::new(
906                "complex_and_or",
907                col("user_id")
908                    .gt_eq(lit(100i64))
909                    .and(col("user_id").lt(lit(200i64)))
910                    .or(col("user_id")
911                        .gt_eq(lit(300i64))
912                        .and(col("user_id").lt(lit(400i64)))),
913                vec![PartitionExpr::new(
914                    Operand::Expr(PartitionExpr::new(
915                        Operand::Expr(PartitionExpr::new(
916                            Operand::Column("user_id".to_string()),
917                            RestrictedOp::GtEq,
918                            Operand::Value(Value::Int64(100)),
919                        )),
920                        RestrictedOp::And,
921                        Operand::Expr(PartitionExpr::new(
922                            Operand::Column("user_id".to_string()),
923                            RestrictedOp::Lt,
924                            Operand::Value(Value::Int64(200)),
925                        )),
926                    )),
927                    RestrictedOp::Or,
928                    Operand::Expr(PartitionExpr::new(
929                        Operand::Expr(PartitionExpr::new(
930                            Operand::Column("user_id".to_string()),
931                            RestrictedOp::GtEq,
932                            Operand::Value(Value::Int64(300)),
933                        )),
934                        RestrictedOp::And,
935                        Operand::Expr(PartitionExpr::new(
936                            Operand::Column("user_id".to_string()),
937                            RestrictedOp::Lt,
938                            Operand::Value(Value::Int64(400)),
939                        )),
940                    )),
941                )],
942                vec!["user_id"],
943            ),
944        ];
945        check_partition_expressions(cases);
946    }
947
948    #[test]
949    fn test_alias_expressions() {
950        let cases = vec![
951            FilterTestCase::new(
952                "simple_alias",
953                col("user_id").alias("uid").eq(lit(100i64)),
954                vec![PartitionExpr::new(
955                    Operand::Column("user_id".to_string()),
956                    RestrictedOp::Eq,
957                    Operand::Value(Value::Int64(100)),
958                )],
959                vec!["user_id"],
960            ),
961            FilterTestCase::new(
962                "nested_alias",
963                col("user_id").alias("uid").alias("u").gt_eq(lit(50i64)),
964                vec![PartitionExpr::new(
965                    Operand::Column("user_id".to_string()),
966                    RestrictedOp::GtEq,
967                    Operand::Value(Value::Int64(50)),
968                )],
969                vec!["user_id"],
970            ),
971            FilterTestCase::new(
972                "complex_alias_with_and_or",
973                col("user_id")
974                    .alias("uid")
975                    .gt_eq(lit(100i64))
976                    .and(col("user_id").alias("u").lt(lit(200i64)))
977                    .or(col("user_id").alias("id").eq(lit(300i64))),
978                vec![PartitionExpr::new(
979                    Operand::Expr(PartitionExpr::new(
980                        Operand::Expr(PartitionExpr::new(
981                            Operand::Column("user_id".to_string()),
982                            RestrictedOp::GtEq,
983                            Operand::Value(Value::Int64(100)),
984                        )),
985                        RestrictedOp::And,
986                        Operand::Expr(PartitionExpr::new(
987                            Operand::Column("user_id".to_string()),
988                            RestrictedOp::Lt,
989                            Operand::Value(Value::Int64(200)),
990                        )),
991                    )),
992                    RestrictedOp::Or,
993                    Operand::Expr(PartitionExpr::new(
994                        Operand::Column("user_id".to_string()),
995                        RestrictedOp::Eq,
996                        Operand::Value(Value::Int64(300)),
997                    )),
998                )],
999                vec!["user_id"],
1000            ),
1001        ];
1002        check_partition_expressions(cases);
1003    }
1004
1005    #[test]
1006    fn test_inlist_expressions() {
1007        let cases = vec![
1008            FilterTestCase::new(
1009                "simple_inlist",
1010                col("user_id").in_list(vec![lit(100i64), lit(200i64), lit(300i64)], false),
1011                vec![PartitionExpr::new(
1012                    Operand::Expr(PartitionExpr::new(
1013                        Operand::Expr(PartitionExpr::new(
1014                            Operand::Column("user_id".to_string()),
1015                            RestrictedOp::Eq,
1016                            Operand::Value(Value::Int64(100)),
1017                        )),
1018                        RestrictedOp::Or,
1019                        Operand::Expr(PartitionExpr::new(
1020                            Operand::Column("user_id".to_string()),
1021                            RestrictedOp::Eq,
1022                            Operand::Value(Value::Int64(200)),
1023                        )),
1024                    )),
1025                    RestrictedOp::Or,
1026                    Operand::Expr(PartitionExpr::new(
1027                        Operand::Column("user_id".to_string()),
1028                        RestrictedOp::Eq,
1029                        Operand::Value(Value::Int64(300)),
1030                    )),
1031                )],
1032                vec!["user_id"],
1033            ),
1034            FilterTestCase::new(
1035                "negated_inlist",
1036                col("user_id").in_list(vec![lit(100i64), lit(200i64)], true),
1037                vec![PartitionExpr::new(
1038                    Operand::Expr(PartitionExpr::new(
1039                        Operand::Column("user_id".to_string()),
1040                        RestrictedOp::NotEq,
1041                        Operand::Value(Value::Int64(100)),
1042                    )),
1043                    RestrictedOp::And,
1044                    Operand::Expr(PartitionExpr::new(
1045                        Operand::Column("user_id".to_string()),
1046                        RestrictedOp::NotEq,
1047                        Operand::Value(Value::Int64(200)),
1048                    )),
1049                )],
1050                vec!["user_id"],
1051            ),
1052            FilterTestCase::new(
1053                "inlist_with_alias",
1054                col("user_id")
1055                    .alias("uid")
1056                    .in_list(vec![lit(100i64), lit(200i64)], false),
1057                vec![PartitionExpr::new(
1058                    Operand::Expr(PartitionExpr::new(
1059                        Operand::Column("user_id".to_string()),
1060                        RestrictedOp::Eq,
1061                        Operand::Value(Value::Int64(100)),
1062                    )),
1063                    RestrictedOp::Or,
1064                    Operand::Expr(PartitionExpr::new(
1065                        Operand::Column("user_id".to_string()),
1066                        RestrictedOp::Eq,
1067                        Operand::Value(Value::Int64(200)),
1068                    )),
1069                )],
1070                vec!["user_id"],
1071            ),
1072        ];
1073        check_partition_expressions(cases);
1074    }
1075
1076    #[test]
1077    fn test_between_expressions() {
1078        let cases = vec![
1079            FilterTestCase::new(
1080                "simple_between",
1081                col("user_id").between(lit(100i64), lit(200i64)),
1082                vec![PartitionExpr::new(
1083                    Operand::Expr(PartitionExpr::new(
1084                        Operand::Column("user_id".to_string()),
1085                        RestrictedOp::GtEq,
1086                        Operand::Value(Value::Int64(100)),
1087                    )),
1088                    RestrictedOp::And,
1089                    Operand::Expr(PartitionExpr::new(
1090                        Operand::Column("user_id".to_string()),
1091                        RestrictedOp::LtEq,
1092                        Operand::Value(Value::Int64(200)),
1093                    )),
1094                )],
1095                vec!["user_id"],
1096            ),
1097            FilterTestCase::new(
1098                "negated_between",
1099                Expr::Between(datafusion_expr::Between {
1100                    expr: Box::new(col("user_id")),
1101                    negated: true,
1102                    low: Box::new(lit(100i64)),
1103                    high: Box::new(lit(200i64)),
1104                }),
1105                vec![PartitionExpr::new(
1106                    Operand::Expr(PartitionExpr::new(
1107                        Operand::Column("user_id".to_string()),
1108                        RestrictedOp::Lt,
1109                        Operand::Value(Value::Int64(100)),
1110                    )),
1111                    RestrictedOp::Or,
1112                    Operand::Expr(PartitionExpr::new(
1113                        Operand::Column("user_id".to_string()),
1114                        RestrictedOp::Gt,
1115                        Operand::Value(Value::Int64(200)),
1116                    )),
1117                )],
1118                vec!["user_id"],
1119            ),
1120            FilterTestCase::new(
1121                "between_with_alias",
1122                col("user_id")
1123                    .alias("uid")
1124                    .between(lit(100i64), lit(200i64)),
1125                vec![PartitionExpr::new(
1126                    Operand::Expr(PartitionExpr::new(
1127                        Operand::Column("user_id".to_string()),
1128                        RestrictedOp::GtEq,
1129                        Operand::Value(Value::Int64(100)),
1130                    )),
1131                    RestrictedOp::And,
1132                    Operand::Expr(PartitionExpr::new(
1133                        Operand::Column("user_id".to_string()),
1134                        RestrictedOp::LtEq,
1135                        Operand::Value(Value::Int64(200)),
1136                    )),
1137                )],
1138                vec!["user_id"],
1139            ),
1140        ];
1141        check_partition_expressions(cases);
1142    }
1143
1144    #[test]
1145    fn test_null_expressions() {
1146        let cases = vec![
1147            FilterTestCase::new(
1148                "is_null",
1149                col("user_id").is_null(),
1150                vec![PartitionExpr::new(
1151                    Operand::Column("user_id".to_string()),
1152                    RestrictedOp::Eq,
1153                    Operand::Value(Value::Null),
1154                )],
1155                vec!["user_id"],
1156            ),
1157            FilterTestCase::new(
1158                "is_not_null",
1159                col("user_id").is_not_null(),
1160                vec![PartitionExpr::new(
1161                    Operand::Column("user_id".to_string()),
1162                    RestrictedOp::NotEq,
1163                    Operand::Value(Value::Null),
1164                )],
1165                vec!["user_id"],
1166            ),
1167            FilterTestCase::new(
1168                "null_with_alias",
1169                col("user_id").alias("uid").is_null(),
1170                vec![PartitionExpr::new(
1171                    Operand::Column("user_id".to_string()),
1172                    RestrictedOp::Eq,
1173                    Operand::Value(Value::Null),
1174                )],
1175                vec!["user_id"],
1176            ),
1177        ];
1178        check_partition_expressions(cases);
1179    }
1180
1181    #[test]
1182    fn test_cast_expressions() {
1183        let cases = vec![
1184            FilterTestCase::new(
1185                "column_cast_requires_type_proof",
1186                Expr::Cast(datafusion_expr::Cast::new(
1187                    Box::new(col("user_id")),
1188                    DataType::Int64,
1189                ))
1190                .eq(lit(100i64)),
1191                vec![],
1192                vec!["user_id"],
1193            ),
1194            FilterTestCase::new(
1195                "cast_with_alias_requires_type_proof",
1196                Expr::Cast(datafusion_expr::Cast::new(
1197                    Box::new(col("user_id").alias("uid")),
1198                    DataType::Int64,
1199                ))
1200                .eq(lit(100i64)),
1201                vec![],
1202                vec!["user_id"],
1203            ),
1204            FilterTestCase::new(
1205                "unsafe_cast",
1206                Expr::Cast(datafusion_expr::Cast::new(
1207                    Box::new(col("user_id")),
1208                    DataType::List(std::sync::Arc::new(
1209                        datafusion::arrow::datatypes::Field::new("item", DataType::Int32, true),
1210                    )),
1211                ))
1212                .eq(lit(100i64)),
1213                vec![],
1214                vec!["user_id"],
1215            ),
1216        ];
1217        check_partition_expressions(cases);
1218    }
1219
1220    #[test]
1221    fn test_not_expressions() {
1222        let cases = vec![
1223            FilterTestCase::new(
1224                "not_equality",
1225                Expr::Not(Box::new(col("user_id").eq(lit(100i64)))),
1226                vec![PartitionExpr::new(
1227                    Operand::Column("user_id".to_string()),
1228                    RestrictedOp::NotEq,
1229                    Operand::Value(Value::Int64(100)),
1230                )],
1231                vec!["user_id"],
1232            ),
1233            FilterTestCase::new(
1234                "not_comparison",
1235                Expr::Not(Box::new(col("user_id").lt(lit(100i64)))),
1236                vec![PartitionExpr::new(
1237                    Operand::Column("user_id".to_string()),
1238                    RestrictedOp::GtEq,
1239                    Operand::Value(Value::Int64(100)),
1240                )],
1241                vec!["user_id"],
1242            ),
1243            FilterTestCase::new(
1244                "not_is_null",
1245                Expr::Not(Box::new(col("user_id").is_null())),
1246                vec![PartitionExpr::new(
1247                    Operand::Column("user_id".to_string()),
1248                    RestrictedOp::NotEq,
1249                    Operand::Value(Value::Null),
1250                )],
1251                vec!["user_id"],
1252            ),
1253            FilterTestCase::new(
1254                "not_with_alias",
1255                Expr::Not(Box::new(col("user_id").alias("uid").eq(lit(100i64)))),
1256                vec![PartitionExpr::new(
1257                    Operand::Column("user_id".to_string()),
1258                    RestrictedOp::NotEq,
1259                    Operand::Value(Value::Int64(100)),
1260                )],
1261                vec!["user_id"],
1262            ),
1263        ];
1264        check_partition_expressions(cases);
1265    }
1266
1267    #[test]
1268    fn test_edge_cases() {
1269        let cases = vec![
1270            FilterTestCase::new(
1271                "qualified_column_name",
1272                {
1273                    let qualified_col = Expr::Column(Column::new(Some("test"), "user_id"));
1274                    qualified_col.eq(lit(100i64))
1275                },
1276                vec![PartitionExpr::new(
1277                    Operand::Column("user_id".to_string()),
1278                    RestrictedOp::Eq,
1279                    Operand::Value(Value::Int64(100)),
1280                )],
1281                vec!["user_id"],
1282            ),
1283            FilterTestCase::new(
1284                "comprehensive_combinations",
1285                {
1286                    let in_expr = col("user_id")
1287                        .alias("uid")
1288                        .in_list(vec![lit(100i64), lit(200i64)], false);
1289                    let between_expr = col("user_id").between(lit(300i64), lit(400i64));
1290                    in_expr.or(between_expr)
1291                },
1292                vec![PartitionExpr::new(
1293                    Operand::Expr(PartitionExpr::new(
1294                        Operand::Expr(PartitionExpr::new(
1295                            Operand::Column("user_id".to_string()),
1296                            RestrictedOp::Eq,
1297                            Operand::Value(Value::Int64(100)),
1298                        )),
1299                        RestrictedOp::Or,
1300                        Operand::Expr(PartitionExpr::new(
1301                            Operand::Column("user_id".to_string()),
1302                            RestrictedOp::Eq,
1303                            Operand::Value(Value::Int64(200)),
1304                        )),
1305                    )),
1306                    RestrictedOp::Or,
1307                    Operand::Expr(PartitionExpr::new(
1308                        Operand::Expr(PartitionExpr::new(
1309                            Operand::Column("user_id".to_string()),
1310                            RestrictedOp::GtEq,
1311                            Operand::Value(Value::Int64(300)),
1312                        )),
1313                        RestrictedOp::And,
1314                        Operand::Expr(PartitionExpr::new(
1315                            Operand::Column("user_id".to_string()),
1316                            RestrictedOp::LtEq,
1317                            Operand::Value(Value::Int64(400)),
1318                        )),
1319                    )),
1320                )],
1321                vec!["user_id"],
1322            ),
1323        ];
1324        check_partition_expressions(cases);
1325    }
1326
1327    #[test]
1328    fn test_mixed_partition_non_partition_expressions() {
1329        let cases = vec![
1330            // Mixed AND expression - should extract only partition part
1331            FilterTestCase::new(
1332                "mixed_and_expression",
1333                col("user_id")
1334                    .eq(lit(100i64))
1335                    .and(col("value").gt(lit(50i64))),
1336                vec![PartitionExpr::new(
1337                    Operand::Column("user_id".to_string()),
1338                    RestrictedOp::Eq,
1339                    Operand::Value(Value::Int64(100)),
1340                )],
1341                vec!["user_id"],
1342            ),
1343            // Mixed OR expression - should be dropped entirely
1344            FilterTestCase::new(
1345                "mixed_or_expression",
1346                col("user_id")
1347                    .between(lit(1i64), lit(10i64))
1348                    .or(col("value").gt(lit(50i64))),
1349                vec![], // Empty result - expression should be dropped
1350                vec!["user_id"],
1351            ),
1352            // Complex mixed AND expression with multiple parts
1353            FilterTestCase::new(
1354                "complex_mixed_and",
1355                col("user_id")
1356                    .gt_eq(lit(100i64))
1357                    .and(col("value").eq(lit(200i64)))
1358                    .and(col("timestamp").lt(lit(1000i64))),
1359                vec![
1360                    PartitionExpr::new(
1361                        Operand::Column("user_id".to_string()),
1362                        RestrictedOp::GtEq,
1363                        Operand::Value(Value::Int64(100)),
1364                    ),
1365                    PartitionExpr::new(
1366                        Operand::Column("timestamp".to_string()),
1367                        RestrictedOp::Lt,
1368                        Operand::Value(Value::Int64(1000)),
1369                    ),
1370                ],
1371                vec!["user_id", "timestamp"], // Both partition columns
1372            ),
1373            // Pure partition expression - should be kept as-is
1374            FilterTestCase::new(
1375                "pure_partition_and",
1376                col("user_id")
1377                    .gt_eq(lit(100i64))
1378                    .and(col("timestamp").lt(lit(1000i64))),
1379                vec![PartitionExpr::new(
1380                    Operand::Expr(PartitionExpr::new(
1381                        Operand::Column("user_id".to_string()),
1382                        RestrictedOp::GtEq,
1383                        Operand::Value(Value::Int64(100)),
1384                    )),
1385                    RestrictedOp::And,
1386                    Operand::Expr(PartitionExpr::new(
1387                        Operand::Column("timestamp".to_string()),
1388                        RestrictedOp::Lt,
1389                        Operand::Value(Value::Int64(1000)),
1390                    )),
1391                )],
1392                vec!["user_id", "timestamp"],
1393            ),
1394            // Pure partition OR expression - should be kept as-is
1395            FilterTestCase::new(
1396                "pure_partition_or",
1397                col("user_id")
1398                    .eq(lit(100i64))
1399                    .or(col("user_id").eq(lit(200i64))),
1400                vec![PartitionExpr::new(
1401                    Operand::Expr(PartitionExpr::new(
1402                        Operand::Column("user_id".to_string()),
1403                        RestrictedOp::Eq,
1404                        Operand::Value(Value::Int64(100)),
1405                    )),
1406                    RestrictedOp::Or,
1407                    Operand::Expr(PartitionExpr::new(
1408                        Operand::Column("user_id".to_string()),
1409                        RestrictedOp::Eq,
1410                        Operand::Value(Value::Int64(200)),
1411                    )),
1412                )],
1413                vec!["user_id"],
1414            ),
1415            // Pure non-partition expression - should be dropped
1416            FilterTestCase::new(
1417                "pure_non_partition",
1418                col("value").gt_eq(lit(100i64)),
1419                vec![], // Empty result - no partition columns involved
1420                vec!["user_id"],
1421            ),
1422            // Complex nested mixed expression
1423            FilterTestCase::new(
1424                "nested_mixed_expression",
1425                (col("user_id")
1426                    .eq(lit(100i64))
1427                    .and(col("value").gt(lit(50i64))))
1428                .or(col("user_id").eq(lit(200i64))),
1429                vec![], // Empty result - OR with mixed sub-expression should be dropped
1430                vec!["user_id"],
1431            ),
1432            // AND with nested OR (mixed) - should extract partition parts only
1433            FilterTestCase::new(
1434                "and_with_nested_mixed_or",
1435                col("user_id")
1436                    .gt_eq(lit(100i64))
1437                    .and(col("value").eq(lit(1i64)).or(col("value").eq(lit(2i64)))),
1438                vec![PartitionExpr::new(
1439                    Operand::Column("user_id".to_string()),
1440                    RestrictedOp::GtEq,
1441                    Operand::Value(Value::Int64(100)),
1442                )],
1443                vec!["user_id"],
1444            ),
1445        ];
1446        check_partition_expressions(cases);
1447    }
1448}
query/dist_plan/predicate_extractor.rs

query/dist_plan/
predicate_extractor.rs
