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

query/dist_plan/
predicate_extractor.rs
