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            _ => false,
210        }
211    }
212}
213// Value range and atomic overlap logic is now refactored into `partition::diff`.
214
215#[cfg(test)]
216mod tests {
217    use datatypes::value::Value;
218    use partition::expr::{Operand, PartitionExpr, RestrictedOp, col};
219    use store_api::storage::RegionId;
220
221    use super::*;
222
223    fn create_test_partition_info(region_id: u64, expr: Option<PartitionExpr>) -> PartitionInfo {
224        PartitionInfo {
225            id: RegionId::new(1, region_id as u32),
226            partition_expr: expr,
227        }
228    }
229
230    #[test]
231    fn test_constraint_pruning_equality() {
232        let partitions = vec![
233            // Region 1: user_id >= 0 AND user_id < 100
234            create_test_partition_info(
235                1,
236                Some(
237                    col("user_id")
238                        .gt_eq(Value::Int64(0))
239                        .and(col("user_id").lt(Value::Int64(100))),
240                ),
241            ),
242            // Region 2: user_id >= 100 AND user_id < 200
243            create_test_partition_info(
244                2,
245                Some(
246                    col("user_id")
247                        .gt_eq(Value::Int64(100))
248                        .and(col("user_id").lt(Value::Int64(200))),
249                ),
250            ),
251            // Region 3: user_id >= 200 AND user_id < 300
252            create_test_partition_info(
253                3,
254                Some(
255                    col("user_id")
256                        .gt_eq(Value::Int64(200))
257                        .and(col("user_id").lt(Value::Int64(300))),
258                ),
259            ),
260        ];
261
262        // Query: user_id = 150 (should only match Region 2)
263        let query_exprs = vec![col("user_id").eq(Value::Int64(150))];
264        let mut column_datatypes = HashMap::default();
265        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
266        let pruned =
267            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
268
269        // Should include Region 2, and potentially others due to conservative approach
270        assert!(pruned.contains(&RegionId::new(1, 2)));
271    }
272
273    #[test]
274    fn test_constraint_pruning_in_list() {
275        let partitions = vec![
276            // Region 1: user_id >= 0 AND user_id < 100
277            create_test_partition_info(
278                1,
279                Some(
280                    col("user_id")
281                        .gt_eq(Value::Int64(0))
282                        .and(col("user_id").lt(Value::Int64(100))),
283                ),
284            ),
285            // Region 2: user_id >= 100 AND user_id < 200
286            create_test_partition_info(
287                2,
288                Some(
289                    col("user_id")
290                        .gt_eq(Value::Int64(100))
291                        .and(col("user_id").lt(Value::Int64(200))),
292                ),
293            ),
294            // Region 3: user_id >= 200 AND user_id < 300
295            create_test_partition_info(
296                3,
297                Some(
298                    col("user_id")
299                        .gt_eq(Value::Int64(200))
300                        .and(col("user_id").lt(Value::Int64(300))),
301                ),
302            ),
303        ];
304
305        // Query: user_id IN (50, 150, 250) - should match all regions
306        let query_exprs = vec![PartitionExpr::new(
307            Operand::Expr(PartitionExpr::new(
308                Operand::Expr(col("user_id").eq(Value::Int64(50))),
309                RestrictedOp::Or,
310                Operand::Expr(col("user_id").eq(Value::Int64(150))),
311            )),
312            RestrictedOp::Or,
313            Operand::Expr(col("user_id").eq(Value::Int64(250))),
314        )];
315
316        let mut column_datatypes = HashMap::default();
317        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
318        let pruned =
319            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
320
321        // Should include regions that can satisfy any of the values
322        assert!(!pruned.is_empty());
323    }
324
325    #[test]
326    fn test_constraint_pruning_range() {
327        let partitions = vec![
328            // Region 1: user_id >= 0 AND user_id < 100
329            create_test_partition_info(
330                1,
331                Some(
332                    col("user_id")
333                        .gt_eq(Value::Int64(0))
334                        .and(col("user_id").lt(Value::Int64(100))),
335                ),
336            ),
337            // Region 2: user_id >= 100 AND user_id < 200
338            create_test_partition_info(
339                2,
340                Some(
341                    col("user_id")
342                        .gt_eq(Value::Int64(100))
343                        .and(col("user_id").lt(Value::Int64(200))),
344                ),
345            ),
346            // Region 3: user_id >= 200 AND user_id < 300
347            create_test_partition_info(
348                3,
349                Some(
350                    col("user_id")
351                        .gt_eq(Value::Int64(200))
352                        .and(col("user_id").lt(Value::Int64(300))),
353                ),
354            ),
355        ];
356
357        // Query: user_id >= 150 (should include regions that can satisfy this constraint)
358        let query_exprs = vec![col("user_id").gt_eq(Value::Int64(150))];
359        let mut column_datatypes = HashMap::default();
360        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
361        let pruned =
362            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
363
364        // With constraint-based approach:
365        // Region 1: [0, 100) - user_id >= 150 is not satisfiable
366        // Region 2: [100, 200) - user_id >= 150 is satisfiable in range [150, 200)
367        // Region 3: [200, 300) - user_id >= 150 is satisfiable (all values >= 200 satisfy user_id >= 150)
368        // Conservative approach may include more regions, but should at least include regions 2 and 3
369        assert!(pruned.len() >= 2);
370        assert!(pruned.contains(&RegionId::new(1, 2))); // Region 2 should be included
371        assert!(pruned.contains(&RegionId::new(1, 3))); // Region 3 should be included
372    }
373
374    #[test]
375    fn test_prune_regions_no_constraints() {
376        let partitions = vec![
377            create_test_partition_info(1, None),
378            create_test_partition_info(2, None),
379        ];
380
381        let constraints = vec![];
382        let column_datatypes = HashMap::default();
383        let pruned =
384            ConstraintPruner::prune_regions(&constraints, &partitions, column_datatypes).unwrap();
385
386        // No constraints should return all regions
387        assert_eq!(pruned.len(), 2);
388    }
389
390    #[test]
391    fn test_missing_partition_expression_returns_all_regions() {
392        let partitions = vec![
393            create_test_partition_info(1, Some(col("user_id").lt(Value::Int64(100)))),
394            create_test_partition_info(2, None),
395            create_test_partition_info(3, Some(col("user_id").gt_eq(Value::Int64(200)))),
396        ];
397        let query_exprs = vec![col("user_id").eq(Value::Int64(150))];
398        let mut column_datatypes = HashMap::default();
399        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
400
401        let pruned =
402            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
403
404        assert_eq!(
405            vec![
406                RegionId::new(1, 1),
407                RegionId::new(1, 2),
408                RegionId::new(1, 3),
409            ],
410            pruned
411        );
412    }
413
414    #[test]
415    fn test_prune_regions_with_simple_equality() {
416        let partitions = vec![
417            // Region 1: user_id >= 0 AND user_id < 100
418            create_test_partition_info(
419                1,
420                Some(
421                    col("user_id")
422                        .gt_eq(Value::Int64(0))
423                        .and(col("user_id").lt(Value::Int64(100))),
424                ),
425            ),
426            // Region 2: user_id >= 100 AND user_id < 200
427            create_test_partition_info(
428                2,
429                Some(
430                    col("user_id")
431                        .gt_eq(Value::Int64(100))
432                        .and(col("user_id").lt(Value::Int64(200))),
433                ),
434            ),
435            // Region 3: user_id >= 200 AND user_id < 300
436            create_test_partition_info(
437                3,
438                Some(
439                    col("user_id")
440                        .gt_eq(Value::Int64(200))
441                        .and(col("user_id").lt(Value::Int64(300))),
442                ),
443            ),
444        ];
445
446        // Query: user_id = 150 (should only match Region 2 which contains values [100, 200))
447        let query_exprs = vec![col("user_id").eq(Value::Int64(150))];
448        let mut column_datatypes = HashMap::default();
449        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
450        let pruned =
451            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
452
453        // user_id = 150 should match Region 2 ([100, 200)) and potentially others due to conservative approach
454        assert!(pruned.contains(&RegionId::new(1, 2)));
455    }
456
457    #[test]
458    fn test_prune_regions_with_or_constraint() {
459        let partitions = vec![
460            // Region 1: user_id >= 0 AND user_id < 100
461            create_test_partition_info(
462                1,
463                Some(
464                    col("user_id")
465                        .gt_eq(Value::Int64(0))
466                        .and(col("user_id").lt(Value::Int64(100))),
467                ),
468            ),
469            // Region 2: user_id >= 100 AND user_id < 200
470            create_test_partition_info(
471                2,
472                Some(
473                    col("user_id")
474                        .gt_eq(Value::Int64(100))
475                        .and(col("user_id").lt(Value::Int64(200))),
476                ),
477            ),
478            // Region 3: user_id >= 200 AND user_id < 300
479            create_test_partition_info(
480                3,
481                Some(
482                    col("user_id")
483                        .gt_eq(Value::Int64(200))
484                        .and(col("user_id").lt(Value::Int64(300))),
485                ),
486            ),
487        ];
488
489        // Query: user_id = 50 OR user_id = 150 OR user_id = 250 - should match all 3 regions
490        let expr1 = col("user_id").eq(Value::Int64(50));
491        let expr2 = col("user_id").eq(Value::Int64(150));
492        let expr3 = col("user_id").eq(Value::Int64(250));
493
494        let or_expr = PartitionExpr::new(
495            Operand::Expr(PartitionExpr::new(
496                Operand::Expr(expr1),
497                RestrictedOp::Or,
498                Operand::Expr(expr2),
499            )),
500            RestrictedOp::Or,
501            Operand::Expr(expr3),
502        );
503
504        let query_exprs = vec![or_expr];
505        let mut column_datatypes = HashMap::default();
506        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
507        let pruned =
508            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
509
510        // Should match all 3 regions: 50 matches Region 1, 150 matches Region 2, 250 matches Region 3
511        assert_eq!(pruned.len(), 3);
512        assert!(pruned.contains(&RegionId::new(1, 1)));
513        assert!(pruned.contains(&RegionId::new(1, 2)));
514        assert!(pruned.contains(&RegionId::new(1, 3)));
515    }
516
517    #[test]
518    fn test_constraint_pruning_no_match() {
519        let partitions = vec![
520            // Region 1: user_id >= 0 AND user_id < 100
521            create_test_partition_info(
522                1,
523                Some(
524                    col("user_id")
525                        .gt_eq(Value::Int64(0))
526                        .and(col("user_id").lt(Value::Int64(100))),
527                ),
528            ),
529            // Region 2: user_id >= 100 AND user_id < 200
530            create_test_partition_info(
531                2,
532                Some(
533                    col("user_id")
534                        .gt_eq(Value::Int64(100))
535                        .and(col("user_id").lt(Value::Int64(200))),
536                ),
537            ),
538        ];
539
540        // Query: user_id = 300 (should match no regions)
541        let query_exprs = vec![col("user_id").eq(Value::Int64(300))];
542        let mut column_datatypes = HashMap::default();
543        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
544        let pruned =
545            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
546
547        // Should match no regions since 300 is outside both partition ranges
548        assert_eq!(pruned.len(), 0);
549    }
550
551    #[test]
552    fn test_constraint_pruning_partial_match() {
553        let partitions = vec![
554            // Region 1: user_id >= 0 AND user_id < 100
555            create_test_partition_info(
556                1,
557                Some(
558                    col("user_id")
559                        .gt_eq(Value::Int64(0))
560                        .and(col("user_id").lt(Value::Int64(100))),
561                ),
562            ),
563            // Region 2: user_id >= 100 AND user_id < 200
564            create_test_partition_info(
565                2,
566                Some(
567                    col("user_id")
568                        .gt_eq(Value::Int64(100))
569                        .and(col("user_id").lt(Value::Int64(200))),
570                ),
571            ),
572        ];
573
574        // Query: user_id >= 50 (should match both regions partially)
575        let query_exprs = vec![col("user_id").gt_eq(Value::Int64(50))];
576        let mut column_datatypes = HashMap::default();
577        column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype());
578        let pruned =
579            ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap();
580
581        // Region 1: [0,100) intersects with [50,∞) -> includes [50,100)
582        // Region 2: [100,200) is fully contained in [50,∞)
583        assert_eq!(pruned.len(), 2);
584        assert!(pruned.contains(&RegionId::new(1, 1)));
585        assert!(pruned.contains(&RegionId::new(1, 2)));
586    }
587}
query/dist_plan/region_pruner.rs

query/dist_plan/
region_pruner.rs
