query/dist_plan/
region_pruner.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//! [`ConstraintPruner`] prunes partition info based on given expressions.
16
17use ahash::{HashMap, HashSet};
18use common_telemetry::debug;
19use datatypes::prelude::ConcreteDataType;
20use datatypes::value::{OrderedFloat, Value};
21use partition::collider::{AtomicExpr, Collider};
22use partition::expr::{Operand, PartitionExpr};
23use partition::manager::PartitionInfo;
24use partition::overlap::atomic_exprs_overlap;
25use store_api::storage::RegionId;
26
27use crate::error::Result;
28
29pub struct ConstraintPruner;
30
31impl ConstraintPruner {
32    /// Prune regions using constraint satisfaction approach
33    ///
34    /// Takes query expressions and partition info, returns matching region IDs
35    pub fn prune_regions(
36        query_expressions: &[PartitionExpr],
37        partitions: &[PartitionInfo],
38        column_datatypes: HashMap<String, ConcreteDataType>,
39    ) -> Result<Vec<RegionId>> {
40        let start = std::time::Instant::now();
41        let all_regions = partitions
42            .iter()
43            .map(|partition| partition.id)
44            .collect::<Vec<_>>();
45        if query_expressions.is_empty() || partitions.is_empty() {
46            // No constraints, return all regions
47            return Ok(all_regions);
48        }
49
50        // Collect all partition expressions for unified normalization
51        let Some(all_partition_expressions) = partitions
52            .iter()
53            .map(|partition| partition.partition_expr.clone())
54            .collect::<Option<Vec<_>>>()
55        else {
56            debug!(
57                "Partition metadata contains a missing partition expression, returning all regions conservatively"
58            );
59            return Ok(all_regions);
60        };
61
62        // Create unified collider with both query and partition expressions for consistent normalization
63        let mut all_expressions = query_expressions.to_vec();
64        all_expressions.extend(all_partition_expressions.iter().cloned());
65        if !Self::normalize_datatype(&mut all_expressions, &column_datatypes) {
66            return Ok(all_regions);
67        }
68
69        let collider = match Collider::new(&all_expressions) {
70            Ok(collider) => collider,
71            Err(err) => {
72                debug!(
73                    "Failed to create unified collider: {}, returning all regions conservatively",
74                    err
75                );
76                return Ok(all_regions);
77            }
78        };
79
80        // Extract query atomic expressions (first N expressions in the collider)
81        let query_atomics: Vec<&AtomicExpr> = collider
82            .atomic_exprs
83            .iter()
84            .filter(|atomic| atomic.source_expr_index < query_expressions.len())
85            .collect();
86
87        let mut candidate_regions = HashSet::default();
88
89        for region_atomics in collider
90            .atomic_exprs
91            .iter()
92            .filter(|atomic| atomic.source_expr_index >= query_expressions.len())
93        {
94            if Self::atomic_sets_overlap(&query_atomics, region_atomics) {
95                let partition_expr_index =
96                    region_atomics.source_expr_index - query_expressions.len();
97                candidate_regions.insert(all_regions[partition_expr_index]);
98            }
99        }
100
101        debug!(
102            "Constraint pruning (cost {}ms): {} -> {} regions",
103            start.elapsed().as_millis(),
104            partitions.len(),
105            candidate_regions.len()
106        );
107
108        Ok(candidate_regions.into_iter().collect())
109    }
110
111    fn atomic_sets_overlap(query_atomics: &[&AtomicExpr], partition_atomic: &AtomicExpr) -> bool {
112        query_atomics
113            .iter()
114            .any(|qa| atomic_exprs_overlap(qa, partition_atomic))
115    }
116
117    fn normalize_datatype(
118        all_expressions: &mut Vec<PartitionExpr>,
119        column_datatypes: &HashMap<String, ConcreteDataType>,
120    ) -> bool {
121        for expr in all_expressions {
122            if !Self::normalize_expr_datatype(&mut expr.lhs, &mut expr.rhs, column_datatypes) {
123                return false;
124            }
125        }
126        true
127    }
128
129    fn normalize_expr_datatype(
130        lhs: &mut Operand,
131        rhs: &mut Operand,
132        column_datatypes: &HashMap<String, ConcreteDataType>,
133    ) -> bool {
134        match (lhs, rhs) {
135            (Operand::Expr(lhs_expr), Operand::Expr(rhs_expr)) => {
136                Self::normalize_expr_datatype(
137                    &mut lhs_expr.lhs,
138                    &mut lhs_expr.rhs,
139                    column_datatypes,
140                ) && Self::normalize_expr_datatype(
141                    &mut rhs_expr.lhs,
142                    &mut rhs_expr.rhs,
143                    column_datatypes,
144                )
145            }
146            (Operand::Column(col_name), Operand::Value(val))
147            | (Operand::Value(val), Operand::Column(col_name)) => {
148                let Some(datatype) = column_datatypes.get(col_name) else {
149                    debug!("Column {} not found from type set, skip pruning", col_name);
150                    return false;
151                };
152
153                match datatype {
154                    ConcreteDataType::Int8(_)
155                    | ConcreteDataType::Int16(_)
156                    | ConcreteDataType::Int32(_)
157                    | ConcreteDataType::Int64(_) => {
158                        let Some(new_lit) = val.as_i64() else {
159                            debug!("Value {:?} cannot be converted to i64", val);
160                            return false;
161                        };
162                        *val = Value::Int64(new_lit);
163                    }
164
165                    ConcreteDataType::UInt8(_)
166                    | ConcreteDataType::UInt16(_)
167                    | ConcreteDataType::UInt32(_)
168                    | ConcreteDataType::UInt64(_) => {
169                        let Some(new_lit) = val.as_u64() else {
170                            debug!("Value {:?} cannot be converted to u64", val);
171                            return false;
172                        };
173                        *val = Value::UInt64(new_lit);
174                    }
175
176                    ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_) => {
177                        let Some(new_lit) = val.as_f64_lossy() else {
178                            debug!("Value {:?} cannot be converted to f64", val);
179                            return false;
180                        };
181
182                        *val = Value::Float64(OrderedFloat(new_lit));
183                    }
184
185                    ConcreteDataType::String(_) | ConcreteDataType::Boolean(_) => {
186                        // no operation needed
187                    }
188
189                    ConcreteDataType::Decimal128(_)
190                    | ConcreteDataType::Binary(_)
191                    | ConcreteDataType::Date(_)
192                    | ConcreteDataType::Timestamp(_)
193                    | ConcreteDataType::Time(_)
194                    | ConcreteDataType::Duration(_)
195                    | ConcreteDataType::Interval(_)
196                    | ConcreteDataType::List(_)
197                    | ConcreteDataType::Dictionary(_)
198                    | ConcreteDataType::Struct(_)
199                    | ConcreteDataType::Json(_)
200                    | ConcreteDataType::Null(_)
201                    | ConcreteDataType::Vector(_) => {
202                        debug!("Unsupported data type {datatype}");
203                        return false;
204                    }
205                }
206
207                true
208            }
209            // Query predicates never constrain a function dimension, so its values all
210            // come from partition rules, already typed as the function output.
211            (Operand::IcebergHash(_), Operand::Value(_))
212            | (Operand::Value(_), Operand::IcebergHash(_)) => true,
213            _ => false,
214        }
215    }
216}
217// Value range and atomic overlap logic is now refactored into `partition::diff`.
218
219#[cfg(test)]
220mod tests {
221    use datatypes::value::Value;
222    use partition::expr::{Operand, PartitionExpr, RestrictedOp, col};
223    use store_api::storage::RegionId;
224
225    use super::*;
226
227    fn create_test_partition_info(region_id: u64, expr: Option<PartitionExpr>) -> PartitionInfo {
228        PartitionInfo {
229            id: RegionId::new(1, region_id as u32),
230            partition_expr: expr,
231        }
232    }
233
234    #[test]
235    fn test_constraint_pruning_equality() {
236        let partitions = vec![
237            // Region 1: user_id >= 0 AND user_id < 100
238            create_test_partition_info(
239                1,
240                Some(
241                    col("user_id")
242                        .gt_eq(Value::Int64(0))
243                        .and(col("user_id").lt(Value::Int64(100))),
244                ),
245            ),
246            // Region 2: user_id >= 100 AND user_id < 200
247            create_test_partition_info(
248                2,
249                Some(
250                    col("user_id")
251                        .gt_eq(Value::Int64(100))
252                        .and(col("user_id").lt(Value::Int64(200))),
253                ),
254            ),
255            // Region 3: user_id >= 200 AND user_id < 300
256            create_test_partition_info(
257                3,
258                Some(
259                    col("user_id")
260                        .gt_eq(Value::Int64(200))
261                        .and(col("user_id").lt(Value::Int64(300))),
262                ),
263            ),
264        ];
265
266        // Query: user_id = 150 (should only match Region 2)
267        let query_exprs = vec![col("user_id").eq(Value::Int64(150))];
268        let mut column_datatypes = HashMap::default();
269        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
270        let pruned =
271            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
272
273        // Should include Region 2, and potentially others due to conservative approach
274        assert!(pruned.contains(&RegionId::new(1, 2)));
275    }
276
277    #[test]
278    fn test_constraint_pruning_with_hash_dimension() {
279        use partition::function::IcebergHash;
280
281        let hash = Operand::IcebergHash(IcebergHash::new("host"));
282        let partitions = [
283            col("ts")
284                .lt(Value::Int64(100))
285                .and(hash.clone().lt(Value::Int32(0))),
286            col("ts")
287                .lt(Value::Int64(100))
288                .and(hash.clone().gt_eq(Value::Int32(0))),
289            col("ts")
290                .gt_eq(Value::Int64(100))
291                .and(hash.clone().lt(Value::Int32(0))),
292            col("ts")
293                .gt_eq(Value::Int64(100))
294                .and(hash.gt_eq(Value::Int32(0))),
295        ]
296        .into_iter()
297        .enumerate()
298        .map(|(i, expr)| create_test_partition_info(i as u64 + 1, Some(expr)))
299        .collect::<Vec<_>>();
300        let column_datatypes = HashMap::from_iter([
301            ("ts".to_string(), ConcreteDataType::int64_datatype()),
302            ("host".to_string(), ConcreteDataType::string_datatype()),
303        ]);
304        let prune = |query: PartitionExpr| {
305            let mut regions =
306                ConstraintPruner::prune_regions(&[query], &partitions, column_datatypes.clone())
307                    .unwrap();
308            regions.sort();
309            regions
310        };
311
312        // A filter on the hashed column can't exclude any hash range.
313        assert_eq!(
314            prune(col("host").eq(Value::String("a".into()))),
315            (1..=4).map(|i| RegionId::new(1, i)).collect::<Vec<_>>()
316        );
317        // Plain columns in the same rule still prune.
318        assert_eq!(
319            prune(
320                col("host")
321                    .eq(Value::String("a".into()))
322                    .and(col("ts").gt_eq(Value::Int64(100)))
323            ),
324            vec![RegionId::new(1, 3), RegionId::new(1, 4)]
325        );
326    }
327
328    #[test]
329    fn test_constraint_pruning_in_list() {
330        let partitions = vec![
331            // Region 1: user_id >= 0 AND user_id < 100
332            create_test_partition_info(
333                1,
334                Some(
335                    col("user_id")
336                        .gt_eq(Value::Int64(0))
337                        .and(col("user_id").lt(Value::Int64(100))),
338                ),
339            ),
340            // Region 2: user_id >= 100 AND user_id < 200
341            create_test_partition_info(
342                2,
343                Some(
344                    col("user_id")
345                        .gt_eq(Value::Int64(100))
346                        .and(col("user_id").lt(Value::Int64(200))),
347                ),
348            ),
349            // Region 3: user_id >= 200 AND user_id < 300
350            create_test_partition_info(
351                3,
352                Some(
353                    col("user_id")
354                        .gt_eq(Value::Int64(200))
355                        .and(col("user_id").lt(Value::Int64(300))),
356                ),
357            ),
358        ];
359
360        // Query: user_id IN (50, 150, 250) - should match all regions
361        let query_exprs = vec![PartitionExpr::new(
362            Operand::Expr(PartitionExpr::new(
363                Operand::Expr(col("user_id").eq(Value::Int64(50))),
364                RestrictedOp::Or,
365                Operand::Expr(col("user_id").eq(Value::Int64(150))),
366            )),
367            RestrictedOp::Or,
368            Operand::Expr(col("user_id").eq(Value::Int64(250))),
369        )];
370
371        let mut column_datatypes = HashMap::default();
372        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
373        let pruned =
374            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
375
376        // Should include regions that can satisfy any of the values
377        assert!(!pruned.is_empty());
378    }
379
380    #[test]
381    fn test_constraint_pruning_range() {
382        let partitions = vec![
383            // Region 1: user_id >= 0 AND user_id < 100
384            create_test_partition_info(
385                1,
386                Some(
387                    col("user_id")
388                        .gt_eq(Value::Int64(0))
389                        .and(col("user_id").lt(Value::Int64(100))),
390                ),
391            ),
392            // Region 2: user_id >= 100 AND user_id < 200
393            create_test_partition_info(
394                2,
395                Some(
396                    col("user_id")
397                        .gt_eq(Value::Int64(100))
398                        .and(col("user_id").lt(Value::Int64(200))),
399                ),
400            ),
401            // Region 3: user_id >= 200 AND user_id < 300
402            create_test_partition_info(
403                3,
404                Some(
405                    col("user_id")
406                        .gt_eq(Value::Int64(200))
407                        .and(col("user_id").lt(Value::Int64(300))),
408                ),
409            ),
410        ];
411
412        // Query: user_id >= 150 (should include regions that can satisfy this constraint)
413        let query_exprs = vec![col("user_id").gt_eq(Value::Int64(150))];
414        let mut column_datatypes = HashMap::default();
415        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
416        let pruned =
417            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
418
419        // With constraint-based approach:
420        // Region 1: [0, 100) - user_id >= 150 is not satisfiable
421        // Region 2: [100, 200) - user_id >= 150 is satisfiable in range [150, 200)
422        // Region 3: [200, 300) - user_id >= 150 is satisfiable (all values >= 200 satisfy user_id >= 150)
423        // Conservative approach may include more regions, but should at least include regions 2 and 3
424        assert!(pruned.len() >= 2);
425        assert!(pruned.contains(&RegionId::new(1, 2))); // Region 2 should be included
426        assert!(pruned.contains(&RegionId::new(1, 3))); // Region 3 should be included
427    }
428
429    #[test]
430    fn test_prune_regions_no_constraints() {
431        let partitions = vec![
432            create_test_partition_info(1, None),
433            create_test_partition_info(2, None),
434        ];
435
436        let constraints = vec![];
437        let column_datatypes = HashMap::default();
438        let pruned =
439            ConstraintPruner::prune_regions(&constraints, &partitions, column_datatypes).unwrap();
440
441        // No constraints should return all regions
442        assert_eq!(pruned.len(), 2);
443    }
444
445    #[test]
446    fn test_missing_partition_expression_returns_all_regions() {
447        let partitions = vec![
448            create_test_partition_info(1, Some(col("user_id").lt(Value::Int64(100)))),
449            create_test_partition_info(2, None),
450            create_test_partition_info(3, Some(col("user_id").gt_eq(Value::Int64(200)))),
451        ];
452        let query_exprs = vec![col("user_id").eq(Value::Int64(150))];
453        let mut column_datatypes = HashMap::default();
454        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
455
456        let pruned =
457            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
458
459        assert_eq!(
460            vec![
461                RegionId::new(1, 1),
462                RegionId::new(1, 2),
463                RegionId::new(1, 3),
464            ],
465            pruned
466        );
467    }
468
469    #[test]
470    fn test_prune_regions_with_simple_equality() {
471        let partitions = vec![
472            // Region 1: user_id >= 0 AND user_id < 100
473            create_test_partition_info(
474                1,
475                Some(
476                    col("user_id")
477                        .gt_eq(Value::Int64(0))
478                        .and(col("user_id").lt(Value::Int64(100))),
479                ),
480            ),
481            // Region 2: user_id >= 100 AND user_id < 200
482            create_test_partition_info(
483                2,
484                Some(
485                    col("user_id")
486                        .gt_eq(Value::Int64(100))
487                        .and(col("user_id").lt(Value::Int64(200))),
488                ),
489            ),
490            // Region 3: user_id >= 200 AND user_id < 300
491            create_test_partition_info(
492                3,
493                Some(
494                    col("user_id")
495                        .gt_eq(Value::Int64(200))
496                        .and(col("user_id").lt(Value::Int64(300))),
497                ),
498            ),
499        ];
500
501        // Query: user_id = 150 (should only match Region 2 which contains values [100, 200))
502        let query_exprs = vec![col("user_id").eq(Value::Int64(150))];
503        let mut column_datatypes = HashMap::default();
504        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
505        let pruned =
506            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
507
508        // user_id = 150 should match Region 2 ([100, 200)) and potentially others due to conservative approach
509        assert!(pruned.contains(&RegionId::new(1, 2)));
510    }
511
512    #[test]
513    fn test_prune_regions_with_or_constraint() {
514        let partitions = vec![
515            // Region 1: user_id >= 0 AND user_id < 100
516            create_test_partition_info(
517                1,
518                Some(
519                    col("user_id")
520                        .gt_eq(Value::Int64(0))
521                        .and(col("user_id").lt(Value::Int64(100))),
522                ),
523            ),
524            // Region 2: user_id >= 100 AND user_id < 200
525            create_test_partition_info(
526                2,
527                Some(
528                    col("user_id")
529                        .gt_eq(Value::Int64(100))
530                        .and(col("user_id").lt(Value::Int64(200))),
531                ),
532            ),
533            // Region 3: user_id >= 200 AND user_id < 300
534            create_test_partition_info(
535                3,
536                Some(
537                    col("user_id")
538                        .gt_eq(Value::Int64(200))
539                        .and(col("user_id").lt(Value::Int64(300))),
540                ),
541            ),
542        ];
543
544        // Query: user_id = 50 OR user_id = 150 OR user_id = 250 - should match all 3 regions
545        let expr1 = col("user_id").eq(Value::Int64(50));
546        let expr2 = col("user_id").eq(Value::Int64(150));
547        let expr3 = col("user_id").eq(Value::Int64(250));
548
549        let or_expr = PartitionExpr::new(
550            Operand::Expr(PartitionExpr::new(
551                Operand::Expr(expr1),
552                RestrictedOp::Or,
553                Operand::Expr(expr2),
554            )),
555            RestrictedOp::Or,
556            Operand::Expr(expr3),
557        );
558
559        let query_exprs = vec![or_expr];
560        let mut column_datatypes = HashMap::default();
561        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
562        let pruned =
563            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
564
565        // Should match all 3 regions: 50 matches Region 1, 150 matches Region 2, 250 matches Region 3
566        assert_eq!(pruned.len(), 3);
567        assert!(pruned.contains(&RegionId::new(1, 1)));
568        assert!(pruned.contains(&RegionId::new(1, 2)));
569        assert!(pruned.contains(&RegionId::new(1, 3)));
570    }
571
572    #[test]
573    fn test_constraint_pruning_no_match() {
574        let partitions = vec![
575            // Region 1: user_id >= 0 AND user_id < 100
576            create_test_partition_info(
577                1,
578                Some(
579                    col("user_id")
580                        .gt_eq(Value::Int64(0))
581                        .and(col("user_id").lt(Value::Int64(100))),
582                ),
583            ),
584            // Region 2: user_id >= 100 AND user_id < 200
585            create_test_partition_info(
586                2,
587                Some(
588                    col("user_id")
589                        .gt_eq(Value::Int64(100))
590                        .and(col("user_id").lt(Value::Int64(200))),
591                ),
592            ),
593        ];
594
595        // Query: user_id = 300 (should match no regions)
596        let query_exprs = vec![col("user_id").eq(Value::Int64(300))];
597        let mut column_datatypes = HashMap::default();
598        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
599        let pruned =
600            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
601
602        // Should match no regions since 300 is outside both partition ranges
603        assert_eq!(pruned.len(), 0);
604    }
605
606    #[test]
607    fn test_constraint_pruning_partial_match() {
608        let partitions = vec![
609            // Region 1: user_id >= 0 AND user_id < 100
610            create_test_partition_info(
611                1,
612                Some(
613                    col("user_id")
614                        .gt_eq(Value::Int64(0))
615                        .and(col("user_id").lt(Value::Int64(100))),
616                ),
617            ),
618            // Region 2: user_id >= 100 AND user_id < 200
619            create_test_partition_info(
620                2,
621                Some(
622                    col("user_id")
623                        .gt_eq(Value::Int64(100))
624                        .and(col("user_id").lt(Value::Int64(200))),
625                ),
626            ),
627        ];
628
629        // Query: user_id >= 50 (should match both regions partially)
630        let query_exprs = vec![col("user_id").gt_eq(Value::Int64(50))];
631        let mut column_datatypes = HashMap::default();
632        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
633        let pruned =
634            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
635
636        // Region 1: [0,100) intersects with [50,∞) -> includes [50,100)
637        // Region 2: [100,200) is fully contained in [50,∞)
638        assert_eq!(pruned.len(), 2);
639        assert!(pruned.contains(&RegionId::new(1, 1)));
640        assert!(pruned.contains(&RegionId::new(1, 2)));
641    }
642}
query/dist_plan/region_pruner.rs

query/dist_plan/
region_pruner.rs
