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(_) => {
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::Value(_) => true, // Values are always safe
256            Operand::Expr(expr) => {
257                Self::expr_only_involves_partition_columns(expr, partition_columns)
258            }
259        }
260    }
261}
262
263/// Converts DataFusion expressions to PartitionExpr
264struct DataFusionExprConverter;
265
266impl DataFusionExprConverter {
267    /// Convert DataFusion Expr to PartitionExpr
268    pub fn convert(expr: &Expr) -> DfResult<PartitionExpr> {
269        match expr {
270            Expr::BinaryExpr(binary_expr) => {
271                let lhs = Self::convert_to_operand(&binary_expr.left)?;
272                let rhs = Self::convert_to_operand(&binary_expr.right)?;
273                let op = Self::convert_operator(&binary_expr.op)?;
274
275                Ok(PartitionExpr::new(lhs, op, rhs))
276            }
277            Expr::InList(inlist_expr) => {
278                // Convert col IN (val1, val2, val3) to col = val1 OR col = val2 OR col = val3
279                // Handle negation: col NOT IN (val1, val2) to col != val1 AND col != val2
280                let column_operand = Self::convert_to_operand(&inlist_expr.expr)?;
281
282                if inlist_expr.list.is_empty() {
283                    return Err(datafusion_common::DataFusionError::Plan(
284                        "InList with empty list is not supported".to_string(),
285                    ));
286                }
287
288                let op = if inlist_expr.negated {
289                    RestrictedOp::NotEq
290                } else {
291                    RestrictedOp::Eq
292                };
293
294                let connector_op = if inlist_expr.negated {
295                    RestrictedOp::And // NOT IN becomes col != val1 AND col != val2
296                } else {
297                    RestrictedOp::Or // IN becomes col = val1 OR col = val2
298                };
299
300                // Convert each value in the list to an equality/inequality expression
301                let mut expressions = Vec::new();
302                for value_expr in &inlist_expr.list {
303                    let value_operand = Self::convert_to_operand(value_expr)?;
304                    expressions.push(PartitionExpr::new(
305                        column_operand.clone(),
306                        op.clone(),
307                        value_operand,
308                    ));
309                }
310
311                // Chain expressions with OR/AND
312                let mut expr_iter = expressions.into_iter();
313                let mut result = expr_iter.next().unwrap();
314                for expr in expr_iter {
315                    result = PartitionExpr::new(
316                        Operand::Expr(result),
317                        connector_op.clone(),
318                        Operand::Expr(expr),
319                    );
320                }
321
322                Ok(result)
323            }
324            Expr::Between(between_expr) => {
325                // Convert col BETWEEN low AND high to col >= low AND col <= high
326                // Handle negation: col NOT BETWEEN low AND high to col < low OR col > high
327                let column_operand = Self::convert_to_operand(&between_expr.expr)?;
328                let low_operand = Self::convert_to_operand(&between_expr.low)?;
329                let high_operand = Self::convert_to_operand(&between_expr.high)?;
330
331                if between_expr.negated {
332                    // NOT BETWEEN: col < low OR col > high
333                    let left_expr =
334                        PartitionExpr::new(column_operand.clone(), RestrictedOp::Lt, low_operand);
335                    let right_expr =
336                        PartitionExpr::new(column_operand, RestrictedOp::Gt, high_operand);
337                    Ok(PartitionExpr::new(
338                        Operand::Expr(left_expr),
339                        RestrictedOp::Or,
340                        Operand::Expr(right_expr),
341                    ))
342                } else {
343                    // BETWEEN: col >= low AND col <= high
344                    let left_expr =
345                        PartitionExpr::new(column_operand.clone(), RestrictedOp::GtEq, low_operand);
346                    let right_expr =
347                        PartitionExpr::new(column_operand, RestrictedOp::LtEq, high_operand);
348                    Ok(PartitionExpr::new(
349                        Operand::Expr(left_expr),
350                        RestrictedOp::And,
351                        Operand::Expr(right_expr),
352                    ))
353                }
354            }
355            Expr::IsNull(expr) => {
356                // Convert col IS NULL to a PartitionExpr
357                let column_operand = Self::convert_to_operand(expr)?;
358                Ok(PartitionExpr::new(
359                    column_operand,
360                    RestrictedOp::Eq,
361                    Operand::Value(Value::Null),
362                ))
363            }
364            Expr::IsNotNull(expr) => {
365                // Convert col IS NOT NULL to a PartitionExpr
366                let column_operand = Self::convert_to_operand(expr)?;
367                Ok(PartitionExpr::new(
368                    column_operand,
369                    RestrictedOp::NotEq,
370                    Operand::Value(Value::Null),
371                ))
372            }
373            Expr::Not(expr) => {
374                // Handle NOT expressions by inverting the inner expression
375                match expr.as_ref() {
376                    Expr::BinaryExpr(binary_expr) => {
377                        let lhs = Self::convert_to_operand(&binary_expr.left)?;
378                        let rhs = Self::convert_to_operand(&binary_expr.right)?;
379                        let inverted_op = Self::invert_operator(&binary_expr.op)?;
380
381                        Ok(PartitionExpr::new(lhs, inverted_op, rhs))
382                    }
383                    Expr::IsNull(inner_expr) => {
384                        // NOT (col IS NULL) becomes col IS NOT NULL
385                        let column_operand = Self::convert_to_operand(inner_expr)?;
386                        Ok(PartitionExpr::new(
387                            column_operand,
388                            RestrictedOp::NotEq,
389                            Operand::Value(Value::Null),
390                        ))
391                    }
392                    Expr::IsNotNull(inner_expr) => {
393                        // NOT (col IS NOT NULL) becomes col IS NULL
394                        let column_operand = Self::convert_to_operand(inner_expr)?;
395                        Ok(PartitionExpr::new(
396                            column_operand,
397                            RestrictedOp::Eq,
398                            Operand::Value(Value::Null),
399                        ))
400                    }
401                    _ => {
402                        debug!(
403                            "Unsupported NOT expression for partition pruning: {:?}",
404                            expr
405                        );
406                        Err(datafusion_common::DataFusionError::Plan(format!(
407                            "NOT expression with inner type {:?} not supported for partition pruning",
408                            expr
409                        )))
410                    }
411                }
412            }
413            _ => Err(datafusion_common::DataFusionError::Plan(format!(
414                "Unsupported expression type for conversion: {:?}",
415                expr
416            ))),
417        }
418    }
419
420    /// Convert DataFusion Expr to Operand
421    fn convert_to_operand(expr: &Expr) -> DfResult<Operand> {
422        match expr {
423            Expr::Column(col) => {
424                // Handle qualified column names (table.column) by extracting just the column name
425                // For partition pruning, we typically only care about the column name itself
426                let column_name = if let Some(relation) = &col.relation {
427                    debug!(
428                        "Using qualified column reference: {}.{}",
429                        relation, col.name
430                    );
431                    col.name.clone()
432                } else {
433                    col.name.clone()
434                };
435                Ok(Operand::Column(column_name))
436            }
437            Expr::Literal(scalar_value, _) => {
438                let value = Value::try_from(scalar_value.clone()).map_err(|err| {
439                    datafusion_common::DataFusionError::Plan(format!(
440                        "Failed to convert literal for partition pruning: {err}"
441                    ))
442                })?;
443                Ok(Operand::Value(value))
444            }
445            Expr::Alias(alias_expr) => {
446                // Unwrap alias to get the actual expression
447                Self::convert_to_operand(&alias_expr.expr)
448            }
449            Expr::Cast(cast_expr) => {
450                // Column casts need input type/value equivalence proof; literals must use the actual cast value.
451                if let Expr::Literal(value, _) = cast_expr.expr.as_ref() {
452                    let value = value.cast_to(cast_expr.field.data_type())?;
453                    let value = Value::try_from(value).map_err(|err| {
454                        datafusion_common::DataFusionError::Plan(format!(
455                            "Failed to convert cast literal for partition pruning: {err}"
456                        ))
457                    })?;
458                    Ok(Operand::Value(value))
459                } else {
460                    debug!("Skipping cast expression for partition pruning: {:?}", expr);
461                    Err(datafusion_common::DataFusionError::Plan(
462                        "Cast expression not supported for partition pruning".to_string(),
463                    ))
464                }
465            }
466            other => {
467                let partition_expr = Self::convert(other)?;
468                Ok(Operand::Expr(partition_expr))
469            }
470        }
471    }
472
473    /// Convert DataFusion Operator to RestrictedOp
474    fn convert_operator(op: &Operator) -> DfResult<RestrictedOp> {
475        match op {
476            Operator::Eq => Ok(RestrictedOp::Eq),
477            Operator::NotEq => Ok(RestrictedOp::NotEq),
478            Operator::Lt => Ok(RestrictedOp::Lt),
479            Operator::LtEq => Ok(RestrictedOp::LtEq),
480            Operator::Gt => Ok(RestrictedOp::Gt),
481            Operator::GtEq => Ok(RestrictedOp::GtEq),
482            Operator::And => Ok(RestrictedOp::And),
483            Operator::Or => Ok(RestrictedOp::Or),
484            _ => Err(datafusion_common::DataFusionError::Plan(format!(
485                "Unsupported operator: {:?}",
486                op
487            ))),
488        }
489    }
490
491    /// Invert a DataFusion Operator for NOT expressions
492    fn invert_operator(op: &Operator) -> DfResult<RestrictedOp> {
493        let Some(negated) = op.negate() else {
494            return Err(datafusion_common::DataFusionError::Plan(format!(
495                "Cannot invert operator: {:?}",
496                op
497            )));
498        };
499        Self::convert_operator(&negated)
500    }
501}
502
503#[cfg(test)]
504mod tests {
505    use std::sync::Arc;
506
507    use datafusion::arrow::datatypes::{DataType, Field, Schema};
508    use datafusion::common::Column;
509    use datafusion::datasource::DefaultTableSource;
510    use datafusion_expr::{LogicalPlanBuilder, col, lit};
511    use datatypes::value::Value;
512    use partition::expr::{Operand, PartitionExpr, RestrictedOp};
513
514    use super::*;
515
516    fn create_test_table_scan() -> LogicalPlan {
517        let schema = Arc::new(Schema::new(vec![
518            Field::new(
519                "timestamp",
520                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
521                false,
522            ),
523            Field::new("user_id", DataType::Int64, false),
524            Field::new("value", DataType::Int64, false),
525        ]));
526
527        let empty_table = datafusion::datasource::empty::EmptyTable::new(schema);
528        let table_source = Arc::new(DefaultTableSource::new(Arc::new(empty_table)));
529
530        LogicalPlanBuilder::scan("test", table_source, None)
531            .unwrap()
532            .build()
533            .unwrap()
534    }
535
536    struct FilterTestCase {
537        name: &'static str,
538        filter_expr: Expr,
539        expected_partition_exprs: Vec<PartitionExpr>,
540        partition_columns: Vec<&'static str>,
541    }
542
543    impl FilterTestCase {
544        fn new(
545            name: &'static str,
546            filter_expr: Expr,
547            expected_partition_exprs: Vec<PartitionExpr>,
548            partition_columns: Vec<&'static str>,
549        ) -> Self {
550            Self {
551                name,
552                filter_expr,
553                expected_partition_exprs,
554                partition_columns,
555            }
556        }
557    }
558
559    /// Helper to check partition expressions for a set of test cases.
560    fn check_partition_expressions(cases: Vec<FilterTestCase>) {
561        for case in cases {
562            let table_scan = create_test_table_scan();
563            let filter = case.filter_expr.clone();
564
565            let plan = LogicalPlanBuilder::from(table_scan)
566                .filter(filter)
567                .unwrap()
568                .build()
569                .unwrap();
570
571            let partition_columns: Vec<String> = case
572                .partition_columns
573                .iter()
574                .map(|s| s.to_string())
575                .collect();
576            let partition_exprs =
577                PredicateExtractor::extract_partition_expressions(&plan, &partition_columns)
578                    .unwrap();
579            let expected = case.expected_partition_exprs.clone();
580            assert_eq!(
581                partition_exprs, expected,
582                "Test case '{}': expected partition expressions {:?}, got {:?}",
583                case.name, expected, partition_exprs
584            );
585        }
586    }
587
588    #[test]
589    fn test_extracts_table_scan_filters() {
590        let table_scan = create_test_table_scan();
591        let filter = col("user_id").gt_eq(lit(100i64));
592        let LogicalPlan::TableScan(scan) = table_scan else {
593            panic!("expected test table scan");
594        };
595        let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan {
596            filters: vec![filter],
597            ..scan
598        });
599
600        let partition_exprs =
601            PredicateExtractor::extract_partition_expressions(&plan, &["user_id".to_string()])
602                .unwrap();
603
604        assert_eq!(
605            partition_exprs,
606            vec![PartitionExpr::new(
607                Operand::Column("user_id".to_string()),
608                RestrictedOp::GtEq,
609                Operand::Value(Value::Int64(100)),
610            )]
611        );
612    }
613
614    #[test]
615    fn test_dictionary_cast_preserves_partition_constraint() {
616        let dictionary_type =
617            DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8));
618        let filter = col("tag").eq(Expr::Cast(datafusion_expr::Cast::new(
619            Box::new(lit("b")),
620            dictionary_type,
621        )));
622
623        let partition_expr = DataFusionExprConverter::convert(&filter).unwrap();
624        assert_eq!(
625            partition_expr,
626            PartitionExpr::new(
627                Operand::Column("tag".to_string()),
628                RestrictedOp::Eq,
629                Operand::Value(Value::String("b".into())),
630            )
631        );
632    }
633
634    #[test]
635    fn test_cast_literals_are_evaluated_and_column_casts_are_skipped() {
636        let casts = [
637            (DataType::Int64, lit(1i64)),
638            (DataType::Float32, lit(1.0f32)),
639            (DataType::Utf8, lit("1")),
640            (DataType::Date32, lit(1i32)),
641            (
642                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Nanosecond, None),
643                lit(1i64),
644            ),
645        ];
646        for (data_type, value) in casts {
647            let cast = Expr::Cast(datafusion_expr::Cast::new(
648                Box::new(col("user_id")),
649                data_type,
650            ));
651            assert!(DataFusionExprConverter::convert(&cast.eq(value)).is_err());
652        }
653
654        let chain = Expr::Cast(datafusion_expr::Cast::new(
655            Box::new(Expr::Cast(datafusion_expr::Cast::new(
656                Box::new(col("user_id")),
657                DataType::Float64,
658            ))),
659            DataType::Int64,
660        ));
661        assert!(DataFusionExprConverter::convert(&chain.eq(lit(1i64))).is_err());
662
663        let cast_int = Expr::Cast(datafusion_expr::Cast::new(
664            Box::new(lit(1.5f64)),
665            DataType::Int64,
666        ));
667        assert_eq!(
668            DataFusionExprConverter::convert(&col("user_id").eq(cast_int))
669                .unwrap()
670                .rhs(),
671            &Operand::Value(Value::Int64(1))
672        );
673
674        let cast_string = Expr::Cast(datafusion_expr::Cast::new(
675            Box::new(lit("12")),
676            DataType::Int64,
677        ));
678        assert_eq!(
679            DataFusionExprConverter::convert(&col("user_id").eq(cast_string))
680                .unwrap()
681                .rhs(),
682            &Operand::Value(Value::Int64(12))
683        );
684
685        let unsupported = Expr::Cast(datafusion_expr::Cast::new(
686            Box::new(lit(1i64)),
687            DataType::Decimal256(10, 2),
688        ));
689        assert!(DataFusionExprConverter::convert_to_operand(&unsupported).is_err());
690    }
691
692    #[test]
693    fn test_cast_predicates_are_skipped_in_and_and_or() {
694        let cast = Expr::Cast(datafusion_expr::Cast::new(
695            Box::new(col("user_id")),
696            DataType::Float32,
697        ));
698        let lower = PartitionExpr::new(
699            Operand::Column("user_id".to_string()),
700            RestrictedOp::GtEq,
701            Operand::Value(Value::Int64(1)),
702        );
703        let upper = PartitionExpr::new(
704            Operand::Column("user_id".to_string()),
705            RestrictedOp::Lt,
706            Operand::Value(Value::Int64(10)),
707        );
708        let bounds = PartitionExpr::new(
709            Operand::Expr(lower.clone()),
710            RestrictedOp::And,
711            Operand::Expr(upper.clone()),
712        );
713        check_partition_expressions(vec![
714            FilterTestCase::new(
715                "cast_in_and_preserves_partition_bound",
716                col("user_id")
717                    .eq(lit(1i64))
718                    .and(cast.clone().eq(lit(1.0f32))),
719                vec![PartitionExpr::new(
720                    Operand::Column("user_id".to_string()),
721                    RestrictedOp::Eq,
722                    Operand::Value(Value::Int64(1)),
723                )],
724                vec!["user_id"],
725            ),
726            FilterTestCase::new(
727                "cast_between_two_bounds_preserves_both",
728                col("user_id")
729                    .gt_eq(lit(1i64))
730                    .and(col("user_id").lt(lit(10i64)))
731                    .and(cast.clone().eq(lit(1.0f32))),
732                vec![bounds.clone()],
733                vec!["user_id"],
734            ),
735            FilterTestCase::new(
736                "cast_left_of_bounds_preserves_both",
737                cast.clone()
738                    .eq(lit(1.0f32))
739                    .and(col("user_id").gt_eq(lit(1i64)))
740                    .and(col("user_id").lt(lit(10i64))),
741                vec![bounds.clone()],
742                vec!["user_id"],
743            ),
744            FilterTestCase::new(
745                "cast_between_bounds_preserves_both",
746                col("user_id")
747                    .gt_eq(lit(1i64))
748                    .and(cast.clone().eq(lit(1.0f32)))
749                    .and(col("user_id").lt(lit(10i64))),
750                vec![bounds],
751                vec!["user_id"],
752            ),
753            FilterTestCase::new(
754                "is_not_null_cast_preserves_partition_bound",
755                col("user_id").eq(lit(1i64)).and(
756                    Expr::Cast(datafusion_expr::Cast::new(
757                        Box::new(col("value")),
758                        DataType::Int32,
759                    ))
760                    .is_not_null(),
761                ),
762                vec![PartitionExpr::new(
763                    Operand::Column("user_id".to_string()),
764                    RestrictedOp::Eq,
765                    Operand::Value(Value::Int64(1)),
766                )],
767                vec!["user_id"],
768            ),
769            FilterTestCase::new(
770                "all_unsupported_conjunction_is_dropped",
771                Expr::Cast(datafusion_expr::Cast::new(
772                    Box::new(col("value")),
773                    DataType::Int32,
774                ))
775                .eq(lit(1i32))
776                .and(
777                    Expr::Cast(datafusion_expr::Cast::new(
778                        Box::new(col("timestamp")),
779                        DataType::Float32,
780                    ))
781                    .eq(lit(1.0f32)),
782                ),
783                vec![],
784                vec!["user_id"],
785            ),
786            FilterTestCase::new(
787                "nested_and_with_or_side_preserves_bound",
788                col("user_id").eq(lit(1i64)).and(
789                    col("user_id")
790                        .eq(lit(2i64))
791                        .or(cast.clone().eq(lit(1.0f32))),
792                ),
793                vec![PartitionExpr::new(
794                    Operand::Column("user_id".to_string()),
795                    RestrictedOp::Eq,
796                    Operand::Value(Value::Int64(1)),
797                )],
798                vec!["user_id"],
799            ),
800            FilterTestCase::new(
801                "or_with_nested_and_cast_drops_whole_expression",
802                col("user_id")
803                    .eq(lit(1i64))
804                    .and(cast.clone().eq(lit(1.0f32)))
805                    .or(col("user_id").eq(lit(2i64))),
806                vec![],
807                vec!["user_id"],
808            ),
809            FilterTestCase::new(
810                "cast_in_or_drops_entire_predicate",
811                col("user_id").eq(lit(1i64)).or(cast.eq(lit(1.0f32))),
812                vec![],
813                vec!["user_id"],
814            ),
815        ]);
816    }
817
818    #[test]
819    fn test_combines_table_scan_filters_as_conjunction() {
820        let table_scan = create_test_table_scan();
821        let filter_a = col("user_id").eq(lit(10i64));
822        let filter_b = col("value").eq(lit(20i64));
823        let LogicalPlan::TableScan(scan) = table_scan else {
824            panic!("expected test table scan");
825        };
826        let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan {
827            filters: vec![filter_a.clone(), filter_b.clone()],
828            ..scan.clone()
829        });
830        let partition_columns = ["user_id".to_string(), "value".to_string()];
831        let expected = vec![PartitionExpr::new(
832            Operand::Expr(PartitionExpr::new(
833                Operand::Column("user_id".to_string()),
834                RestrictedOp::Eq,
835                Operand::Value(Value::Int64(10)),
836            )),
837            RestrictedOp::And,
838            Operand::Expr(PartitionExpr::new(
839                Operand::Column("value".to_string()),
840                RestrictedOp::Eq,
841                Operand::Value(Value::Int64(20)),
842            )),
843        )];
844
845        let partition_exprs =
846            PredicateExtractor::extract_partition_expressions(&plan, &partition_columns).unwrap();
847        assert_eq!(partition_exprs, expected);
848
849        let unsupported = Expr::Cast(datafusion_expr::Cast::new(
850            Box::new(col("timestamp")),
851            DataType::Float32,
852        ))
853        .eq(lit(1.0f32));
854        let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan {
855            filters: vec![filter_a, unsupported, filter_b],
856            ..scan
857        });
858
859        let partition_exprs =
860            PredicateExtractor::extract_partition_expressions(&plan, &partition_columns).unwrap();
861
862        assert_eq!(partition_exprs, expected);
863    }
864
865    #[test]
866    fn test_basic_constraints_extraction() {
867        let cases = vec![
868            FilterTestCase::new(
869                "non_partition_column_ignored",
870                col("value").gt_eq(lit(100i64)),
871                vec![],
872                vec!["user_id"],
873            ),
874            FilterTestCase::new(
875                "simple_constraint",
876                col("user_id").gt_eq(lit(100i64)),
877                vec![PartitionExpr::new(
878                    Operand::Column("user_id".to_string()),
879                    RestrictedOp::GtEq,
880                    Operand::Value(Value::Int64(100)),
881                )],
882                vec!["user_id"],
883            ),
884            FilterTestCase::new(
885                "or_expression",
886                col("user_id")
887                    .eq(lit(100i64))
888                    .or(col("user_id").eq(lit(200i64))),
889                vec![PartitionExpr::new(
890                    Operand::Expr(PartitionExpr::new(
891                        Operand::Column("user_id".to_string()),
892                        RestrictedOp::Eq,
893                        Operand::Value(Value::Int64(100)),
894                    )),
895                    RestrictedOp::Or,
896                    Operand::Expr(PartitionExpr::new(
897                        Operand::Column("user_id".to_string()),
898                        RestrictedOp::Eq,
899                        Operand::Value(Value::Int64(200)),
900                    )),
901                )],
902                vec!["user_id"],
903            ),
904            FilterTestCase::new(
905                "complex_and_or",
906                col("user_id")
907                    .gt_eq(lit(100i64))
908                    .and(col("user_id").lt(lit(200i64)))
909                    .or(col("user_id")
910                        .gt_eq(lit(300i64))
911                        .and(col("user_id").lt(lit(400i64)))),
912                vec![PartitionExpr::new(
913                    Operand::Expr(PartitionExpr::new(
914                        Operand::Expr(PartitionExpr::new(
915                            Operand::Column("user_id".to_string()),
916                            RestrictedOp::GtEq,
917                            Operand::Value(Value::Int64(100)),
918                        )),
919                        RestrictedOp::And,
920                        Operand::Expr(PartitionExpr::new(
921                            Operand::Column("user_id".to_string()),
922                            RestrictedOp::Lt,
923                            Operand::Value(Value::Int64(200)),
924                        )),
925                    )),
926                    RestrictedOp::Or,
927                    Operand::Expr(PartitionExpr::new(
928                        Operand::Expr(PartitionExpr::new(
929                            Operand::Column("user_id".to_string()),
930                            RestrictedOp::GtEq,
931                            Operand::Value(Value::Int64(300)),
932                        )),
933                        RestrictedOp::And,
934                        Operand::Expr(PartitionExpr::new(
935                            Operand::Column("user_id".to_string()),
936                            RestrictedOp::Lt,
937                            Operand::Value(Value::Int64(400)),
938                        )),
939                    )),
940                )],
941                vec!["user_id"],
942            ),
943        ];
944        check_partition_expressions(cases);
945    }
946
947    #[test]
948    fn test_alias_expressions() {
949        let cases = vec![
950            FilterTestCase::new(
951                "simple_alias",
952                col("user_id").alias("uid").eq(lit(100i64)),
953                vec![PartitionExpr::new(
954                    Operand::Column("user_id".to_string()),
955                    RestrictedOp::Eq,
956                    Operand::Value(Value::Int64(100)),
957                )],
958                vec!["user_id"],
959            ),
960            FilterTestCase::new(
961                "nested_alias",
962                col("user_id").alias("uid").alias("u").gt_eq(lit(50i64)),
963                vec![PartitionExpr::new(
964                    Operand::Column("user_id".to_string()),
965                    RestrictedOp::GtEq,
966                    Operand::Value(Value::Int64(50)),
967                )],
968                vec!["user_id"],
969            ),
970            FilterTestCase::new(
971                "complex_alias_with_and_or",
972                col("user_id")
973                    .alias("uid")
974                    .gt_eq(lit(100i64))
975                    .and(col("user_id").alias("u").lt(lit(200i64)))
976                    .or(col("user_id").alias("id").eq(lit(300i64))),
977                vec![PartitionExpr::new(
978                    Operand::Expr(PartitionExpr::new(
979                        Operand::Expr(PartitionExpr::new(
980                            Operand::Column("user_id".to_string()),
981                            RestrictedOp::GtEq,
982                            Operand::Value(Value::Int64(100)),
983                        )),
984                        RestrictedOp::And,
985                        Operand::Expr(PartitionExpr::new(
986                            Operand::Column("user_id".to_string()),
987                            RestrictedOp::Lt,
988                            Operand::Value(Value::Int64(200)),
989                        )),
990                    )),
991                    RestrictedOp::Or,
992                    Operand::Expr(PartitionExpr::new(
993                        Operand::Column("user_id".to_string()),
994                        RestrictedOp::Eq,
995                        Operand::Value(Value::Int64(300)),
996                    )),
997                )],
998                vec!["user_id"],
999            ),
1000        ];
1001        check_partition_expressions(cases);
1002    }
1003
1004    #[test]
1005    fn test_inlist_expressions() {
1006        let cases = vec![
1007            FilterTestCase::new(
1008                "simple_inlist",
1009                col("user_id").in_list(vec![lit(100i64), lit(200i64), lit(300i64)], false),
1010                vec![PartitionExpr::new(
1011                    Operand::Expr(PartitionExpr::new(
1012                        Operand::Expr(PartitionExpr::new(
1013                            Operand::Column("user_id".to_string()),
1014                            RestrictedOp::Eq,
1015                            Operand::Value(Value::Int64(100)),
1016                        )),
1017                        RestrictedOp::Or,
1018                        Operand::Expr(PartitionExpr::new(
1019                            Operand::Column("user_id".to_string()),
1020                            RestrictedOp::Eq,
1021                            Operand::Value(Value::Int64(200)),
1022                        )),
1023                    )),
1024                    RestrictedOp::Or,
1025                    Operand::Expr(PartitionExpr::new(
1026                        Operand::Column("user_id".to_string()),
1027                        RestrictedOp::Eq,
1028                        Operand::Value(Value::Int64(300)),
1029                    )),
1030                )],
1031                vec!["user_id"],
1032            ),
1033            FilterTestCase::new(
1034                "negated_inlist",
1035                col("user_id").in_list(vec![lit(100i64), lit(200i64)], true),
1036                vec![PartitionExpr::new(
1037                    Operand::Expr(PartitionExpr::new(
1038                        Operand::Column("user_id".to_string()),
1039                        RestrictedOp::NotEq,
1040                        Operand::Value(Value::Int64(100)),
1041                    )),
1042                    RestrictedOp::And,
1043                    Operand::Expr(PartitionExpr::new(
1044                        Operand::Column("user_id".to_string()),
1045                        RestrictedOp::NotEq,
1046                        Operand::Value(Value::Int64(200)),
1047                    )),
1048                )],
1049                vec!["user_id"],
1050            ),
1051            FilterTestCase::new(
1052                "inlist_with_alias",
1053                col("user_id")
1054                    .alias("uid")
1055                    .in_list(vec![lit(100i64), lit(200i64)], false),
1056                vec![PartitionExpr::new(
1057                    Operand::Expr(PartitionExpr::new(
1058                        Operand::Column("user_id".to_string()),
1059                        RestrictedOp::Eq,
1060                        Operand::Value(Value::Int64(100)),
1061                    )),
1062                    RestrictedOp::Or,
1063                    Operand::Expr(PartitionExpr::new(
1064                        Operand::Column("user_id".to_string()),
1065                        RestrictedOp::Eq,
1066                        Operand::Value(Value::Int64(200)),
1067                    )),
1068                )],
1069                vec!["user_id"],
1070            ),
1071        ];
1072        check_partition_expressions(cases);
1073    }
1074
1075    #[test]
1076    fn test_between_expressions() {
1077        let cases = vec![
1078            FilterTestCase::new(
1079                "simple_between",
1080                col("user_id").between(lit(100i64), lit(200i64)),
1081                vec![PartitionExpr::new(
1082                    Operand::Expr(PartitionExpr::new(
1083                        Operand::Column("user_id".to_string()),
1084                        RestrictedOp::GtEq,
1085                        Operand::Value(Value::Int64(100)),
1086                    )),
1087                    RestrictedOp::And,
1088                    Operand::Expr(PartitionExpr::new(
1089                        Operand::Column("user_id".to_string()),
1090                        RestrictedOp::LtEq,
1091                        Operand::Value(Value::Int64(200)),
1092                    )),
1093                )],
1094                vec!["user_id"],
1095            ),
1096            FilterTestCase::new(
1097                "negated_between",
1098                Expr::Between(datafusion_expr::Between {
1099                    expr: Box::new(col("user_id")),
1100                    negated: true,
1101                    low: Box::new(lit(100i64)),
1102                    high: Box::new(lit(200i64)),
1103                }),
1104                vec![PartitionExpr::new(
1105                    Operand::Expr(PartitionExpr::new(
1106                        Operand::Column("user_id".to_string()),
1107                        RestrictedOp::Lt,
1108                        Operand::Value(Value::Int64(100)),
1109                    )),
1110                    RestrictedOp::Or,
1111                    Operand::Expr(PartitionExpr::new(
1112                        Operand::Column("user_id".to_string()),
1113                        RestrictedOp::Gt,
1114                        Operand::Value(Value::Int64(200)),
1115                    )),
1116                )],
1117                vec!["user_id"],
1118            ),
1119            FilterTestCase::new(
1120                "between_with_alias",
1121                col("user_id")
1122                    .alias("uid")
1123                    .between(lit(100i64), lit(200i64)),
1124                vec![PartitionExpr::new(
1125                    Operand::Expr(PartitionExpr::new(
1126                        Operand::Column("user_id".to_string()),
1127                        RestrictedOp::GtEq,
1128                        Operand::Value(Value::Int64(100)),
1129                    )),
1130                    RestrictedOp::And,
1131                    Operand::Expr(PartitionExpr::new(
1132                        Operand::Column("user_id".to_string()),
1133                        RestrictedOp::LtEq,
1134                        Operand::Value(Value::Int64(200)),
1135                    )),
1136                )],
1137                vec!["user_id"],
1138            ),
1139        ];
1140        check_partition_expressions(cases);
1141    }
1142
1143    #[test]
1144    fn test_null_expressions() {
1145        let cases = vec![
1146            FilterTestCase::new(
1147                "is_null",
1148                col("user_id").is_null(),
1149                vec![PartitionExpr::new(
1150                    Operand::Column("user_id".to_string()),
1151                    RestrictedOp::Eq,
1152                    Operand::Value(Value::Null),
1153                )],
1154                vec!["user_id"],
1155            ),
1156            FilterTestCase::new(
1157                "is_not_null",
1158                col("user_id").is_not_null(),
1159                vec![PartitionExpr::new(
1160                    Operand::Column("user_id".to_string()),
1161                    RestrictedOp::NotEq,
1162                    Operand::Value(Value::Null),
1163                )],
1164                vec!["user_id"],
1165            ),
1166            FilterTestCase::new(
1167                "null_with_alias",
1168                col("user_id").alias("uid").is_null(),
1169                vec![PartitionExpr::new(
1170                    Operand::Column("user_id".to_string()),
1171                    RestrictedOp::Eq,
1172                    Operand::Value(Value::Null),
1173                )],
1174                vec!["user_id"],
1175            ),
1176        ];
1177        check_partition_expressions(cases);
1178    }
1179
1180    #[test]
1181    fn test_cast_expressions() {
1182        let cases = vec![
1183            FilterTestCase::new(
1184                "column_cast_requires_type_proof",
1185                Expr::Cast(datafusion_expr::Cast::new(
1186                    Box::new(col("user_id")),
1187                    DataType::Int64,
1188                ))
1189                .eq(lit(100i64)),
1190                vec![],
1191                vec!["user_id"],
1192            ),
1193            FilterTestCase::new(
1194                "cast_with_alias_requires_type_proof",
1195                Expr::Cast(datafusion_expr::Cast::new(
1196                    Box::new(col("user_id").alias("uid")),
1197                    DataType::Int64,
1198                ))
1199                .eq(lit(100i64)),
1200                vec![],
1201                vec!["user_id"],
1202            ),
1203            FilterTestCase::new(
1204                "unsafe_cast",
1205                Expr::Cast(datafusion_expr::Cast::new(
1206                    Box::new(col("user_id")),
1207                    DataType::List(std::sync::Arc::new(
1208                        datafusion::arrow::datatypes::Field::new("item", DataType::Int32, true),
1209                    )),
1210                ))
1211                .eq(lit(100i64)),
1212                vec![],
1213                vec!["user_id"],
1214            ),
1215        ];
1216        check_partition_expressions(cases);
1217    }
1218
1219    #[test]
1220    fn test_not_expressions() {
1221        let cases = vec![
1222            FilterTestCase::new(
1223                "not_equality",
1224                Expr::Not(Box::new(col("user_id").eq(lit(100i64)))),
1225                vec![PartitionExpr::new(
1226                    Operand::Column("user_id".to_string()),
1227                    RestrictedOp::NotEq,
1228                    Operand::Value(Value::Int64(100)),
1229                )],
1230                vec!["user_id"],
1231            ),
1232            FilterTestCase::new(
1233                "not_comparison",
1234                Expr::Not(Box::new(col("user_id").lt(lit(100i64)))),
1235                vec![PartitionExpr::new(
1236                    Operand::Column("user_id".to_string()),
1237                    RestrictedOp::GtEq,
1238                    Operand::Value(Value::Int64(100)),
1239                )],
1240                vec!["user_id"],
1241            ),
1242            FilterTestCase::new(
1243                "not_is_null",
1244                Expr::Not(Box::new(col("user_id").is_null())),
1245                vec![PartitionExpr::new(
1246                    Operand::Column("user_id".to_string()),
1247                    RestrictedOp::NotEq,
1248                    Operand::Value(Value::Null),
1249                )],
1250                vec!["user_id"],
1251            ),
1252            FilterTestCase::new(
1253                "not_with_alias",
1254                Expr::Not(Box::new(col("user_id").alias("uid").eq(lit(100i64)))),
1255                vec![PartitionExpr::new(
1256                    Operand::Column("user_id".to_string()),
1257                    RestrictedOp::NotEq,
1258                    Operand::Value(Value::Int64(100)),
1259                )],
1260                vec!["user_id"],
1261            ),
1262        ];
1263        check_partition_expressions(cases);
1264    }
1265
1266    #[test]
1267    fn test_edge_cases() {
1268        let cases = vec![
1269            FilterTestCase::new(
1270                "qualified_column_name",
1271                {
1272                    let qualified_col = Expr::Column(Column::new(Some("test"), "user_id"));
1273                    qualified_col.eq(lit(100i64))
1274                },
1275                vec![PartitionExpr::new(
1276                    Operand::Column("user_id".to_string()),
1277                    RestrictedOp::Eq,
1278                    Operand::Value(Value::Int64(100)),
1279                )],
1280                vec!["user_id"],
1281            ),
1282            FilterTestCase::new(
1283                "comprehensive_combinations",
1284                {
1285                    let in_expr = col("user_id")
1286                        .alias("uid")
1287                        .in_list(vec![lit(100i64), lit(200i64)], false);
1288                    let between_expr = col("user_id").between(lit(300i64), lit(400i64));
1289                    in_expr.or(between_expr)
1290                },
1291                vec![PartitionExpr::new(
1292                    Operand::Expr(PartitionExpr::new(
1293                        Operand::Expr(PartitionExpr::new(
1294                            Operand::Column("user_id".to_string()),
1295                            RestrictedOp::Eq,
1296                            Operand::Value(Value::Int64(100)),
1297                        )),
1298                        RestrictedOp::Or,
1299                        Operand::Expr(PartitionExpr::new(
1300                            Operand::Column("user_id".to_string()),
1301                            RestrictedOp::Eq,
1302                            Operand::Value(Value::Int64(200)),
1303                        )),
1304                    )),
1305                    RestrictedOp::Or,
1306                    Operand::Expr(PartitionExpr::new(
1307                        Operand::Expr(PartitionExpr::new(
1308                            Operand::Column("user_id".to_string()),
1309                            RestrictedOp::GtEq,
1310                            Operand::Value(Value::Int64(300)),
1311                        )),
1312                        RestrictedOp::And,
1313                        Operand::Expr(PartitionExpr::new(
1314                            Operand::Column("user_id".to_string()),
1315                            RestrictedOp::LtEq,
1316                            Operand::Value(Value::Int64(400)),
1317                        )),
1318                    )),
1319                )],
1320                vec!["user_id"],
1321            ),
1322        ];
1323        check_partition_expressions(cases);
1324    }
1325
1326    #[test]
1327    fn test_mixed_partition_non_partition_expressions() {
1328        let cases = vec![
1329            // Mixed AND expression - should extract only partition part
1330            FilterTestCase::new(
1331                "mixed_and_expression",
1332                col("user_id")
1333                    .eq(lit(100i64))
1334                    .and(col("value").gt(lit(50i64))),
1335                vec![PartitionExpr::new(
1336                    Operand::Column("user_id".to_string()),
1337                    RestrictedOp::Eq,
1338                    Operand::Value(Value::Int64(100)),
1339                )],
1340                vec!["user_id"],
1341            ),
1342            // Mixed OR expression - should be dropped entirely
1343            FilterTestCase::new(
1344                "mixed_or_expression",
1345                col("user_id")
1346                    .between(lit(1i64), lit(10i64))
1347                    .or(col("value").gt(lit(50i64))),
1348                vec![], // Empty result - expression should be dropped
1349                vec!["user_id"],
1350            ),
1351            // Complex mixed AND expression with multiple parts
1352            FilterTestCase::new(
1353                "complex_mixed_and",
1354                col("user_id")
1355                    .gt_eq(lit(100i64))
1356                    .and(col("value").eq(lit(200i64)))
1357                    .and(col("timestamp").lt(lit(1000i64))),
1358                vec![
1359                    PartitionExpr::new(
1360                        Operand::Column("user_id".to_string()),
1361                        RestrictedOp::GtEq,
1362                        Operand::Value(Value::Int64(100)),
1363                    ),
1364                    PartitionExpr::new(
1365                        Operand::Column("timestamp".to_string()),
1366                        RestrictedOp::Lt,
1367                        Operand::Value(Value::Int64(1000)),
1368                    ),
1369                ],
1370                vec!["user_id", "timestamp"], // Both partition columns
1371            ),
1372            // Pure partition expression - should be kept as-is
1373            FilterTestCase::new(
1374                "pure_partition_and",
1375                col("user_id")
1376                    .gt_eq(lit(100i64))
1377                    .and(col("timestamp").lt(lit(1000i64))),
1378                vec![PartitionExpr::new(
1379                    Operand::Expr(PartitionExpr::new(
1380                        Operand::Column("user_id".to_string()),
1381                        RestrictedOp::GtEq,
1382                        Operand::Value(Value::Int64(100)),
1383                    )),
1384                    RestrictedOp::And,
1385                    Operand::Expr(PartitionExpr::new(
1386                        Operand::Column("timestamp".to_string()),
1387                        RestrictedOp::Lt,
1388                        Operand::Value(Value::Int64(1000)),
1389                    )),
1390                )],
1391                vec!["user_id", "timestamp"],
1392            ),
1393            // Pure partition OR expression - should be kept as-is
1394            FilterTestCase::new(
1395                "pure_partition_or",
1396                col("user_id")
1397                    .eq(lit(100i64))
1398                    .or(col("user_id").eq(lit(200i64))),
1399                vec![PartitionExpr::new(
1400                    Operand::Expr(PartitionExpr::new(
1401                        Operand::Column("user_id".to_string()),
1402                        RestrictedOp::Eq,
1403                        Operand::Value(Value::Int64(100)),
1404                    )),
1405                    RestrictedOp::Or,
1406                    Operand::Expr(PartitionExpr::new(
1407                        Operand::Column("user_id".to_string()),
1408                        RestrictedOp::Eq,
1409                        Operand::Value(Value::Int64(200)),
1410                    )),
1411                )],
1412                vec!["user_id"],
1413            ),
1414            // Pure non-partition expression - should be dropped
1415            FilterTestCase::new(
1416                "pure_non_partition",
1417                col("value").gt_eq(lit(100i64)),
1418                vec![], // Empty result - no partition columns involved
1419                vec!["user_id"],
1420            ),
1421            // Complex nested mixed expression
1422            FilterTestCase::new(
1423                "nested_mixed_expression",
1424                (col("user_id")
1425                    .eq(lit(100i64))
1426                    .and(col("value").gt(lit(50i64))))
1427                .or(col("user_id").eq(lit(200i64))),
1428                vec![], // Empty result - OR with mixed sub-expression should be dropped
1429                vec!["user_id"],
1430            ),
1431            // AND with nested OR (mixed) - should extract partition parts only
1432            FilterTestCase::new(
1433                "and_with_nested_mixed_or",
1434                col("user_id")
1435                    .gt_eq(lit(100i64))
1436                    .and(col("value").eq(lit(1i64)).or(col("value").eq(lit(2i64)))),
1437                vec![PartitionExpr::new(
1438                    Operand::Column("user_id".to_string()),
1439                    RestrictedOp::GtEq,
1440                    Operand::Value(Value::Int64(100)),
1441                )],
1442                vec!["user_id"],
1443            ),
1444        ];
1445        check_partition_expressions(cases);
1446    }
1447}
query/dist_plan/predicate_extractor.rs

query/dist_plan/
predicate_extractor.rs
