partition/
checker.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
15use std::collections::{BTreeMap, HashMap};
16use std::sync::Arc;
17
18use datatypes::arrow::array::{BooleanArray, Float64Array, Float64Builder, RecordBatch};
19use datatypes::arrow::datatypes::{DataType, Field, Schema};
20use datatypes::value::OrderedF64;
21
22use crate::collider::{CHECK_STEP, Collider, NORMALIZE_STEP, PartitionDim};
23use crate::error::{
24    CheckpointNotCoveredSnafu, CheckpointOverlappedSnafu, DuplicateExprSnafu, Result,
25};
26use crate::expr::{PartitionExpr, RestrictedOp};
27use crate::multi_dim::MultiDimPartitionRule;
28
29pub struct PartitionChecker<'a> {
30    rule: &'a MultiDimPartitionRule,
31    collider: Collider<'a>,
32    /// Dimensions in the matrix. A dimension's field name in the matrix is its
33    /// index here, since a column may be named like a function call.
34    dims: Vec<PartitionDim>,
35}
36
37impl<'a> PartitionChecker<'a> {
38    pub fn try_new(rule: &'a MultiDimPartitionRule) -> Result<Self> {
39        let collider = Collider::new(rule.exprs())?;
40        let mut dims = collider
41            .normalized_values
42            .iter()
43            .filter(|(_, values)| !values.is_empty())
44            .map(|(dim, _)| dim.clone())
45            .collect::<Vec<_>>();
46        dims.sort_unstable();
47        Ok(Self {
48            rule,
49            collider,
50            dims,
51        })
52    }
53
54    pub fn check(&self) -> Result<()> {
55        self.run()?;
56        Ok(())
57    }
58}
59
60// Logic of checking rules
61impl<'a> PartitionChecker<'a> {
62    fn run(&self) -> Result<()> {
63        // Sort atomic exprs and check uniqueness
64        let mut atomic_exprs = BTreeMap::new();
65        for expr in self.collider.atomic_exprs.iter() {
66            let key = &expr.nucleons;
67            atomic_exprs.insert(key, expr);
68        }
69        if atomic_exprs.len() != self.collider.atomic_exprs.len() {
70            // Find the duplication for error message
71            for expr in self.collider.atomic_exprs.iter() {
72                if atomic_exprs.get(&expr.nucleons).unwrap().source_expr_index
73                    != expr.source_expr_index
74                {
75                    let expr = self.rule.exprs()[expr.source_expr_index].clone();
76                    return DuplicateExprSnafu { expr }.fail();
77                }
78            }
79            // Or return a placeholder. This should never happen.
80            return DuplicateExprSnafu {
81                expr: PartitionExpr::new(
82                    crate::expr::Operand::Column("unknown".to_string()),
83                    RestrictedOp::Eq,
84                    crate::expr::Operand::Column("expr".to_string()),
85                ),
86            }
87            .fail();
88        }
89
90        // TODO(ruihang): merge atomic exprs to improve checker's performance
91
92        // matrix test
93        let field_names = (0..self.dims.len())
94            .map(|index| index.to_string())
95            .collect::<Vec<_>>();
96        let mut matrix_foundation = HashMap::new();
97        for (dim, field_name) in self.dims.iter().zip(&field_names) {
98            let values = &self.collider.normalized_values[dim];
99            let mut cornerstones = Vec::with_capacity(values.len() * 2 + 1);
100            cornerstones.push(values[0].1 - CHECK_STEP);
101            for value in values {
102                cornerstones.push(value.1);
103                cornerstones.push(value.1 + CHECK_STEP);
104            }
105            matrix_foundation.insert(field_name.as_str(), cornerstones);
106        }
107
108        // If there are no values, the rule is empty and valid.
109        if matrix_foundation.is_empty() {
110            return Ok(());
111        }
112
113        let matrix_generator = MatrixGenerator::new(matrix_foundation);
114
115        // Process data in batches using iterator
116        let mut results = Vec::with_capacity(self.collider.atomic_exprs.len());
117        let physical_exprs = self
118            .collider
119            .atomic_exprs
120            .iter()
121            .map(|expr| {
122                expr.to_physical_expr(matrix_generator.schema(), |dim| {
123                    // Every nucleon's dimension has a value, so it is in `dims`.
124                    self.dims.binary_search(dim).unwrap().to_string()
125                })
126            })
127            .collect::<Vec<_>>();
128        for batch in matrix_generator {
129            results.clear();
130            for physical_expr in &physical_exprs {
131                let columnar_result = physical_expr.evaluate(&batch).unwrap();
132                let array_result = columnar_result.into_array(batch.num_rows()).unwrap();
133                results.push(array_result);
134            }
135            let boolean_results = results
136                .iter()
137                .map(|result| result.as_any().downcast_ref::<BooleanArray>().unwrap())
138                .collect::<Vec<_>>();
139
140            // sum and check results for this batch
141            for i in 0..batch.num_rows() {
142                let mut true_count = 0;
143                for result in boolean_results.iter() {
144                    if result.value(i) {
145                        true_count += 1;
146                    }
147                }
148
149                if true_count == 0 {
150                    return CheckpointNotCoveredSnafu {
151                        checkpoint: self.remap_checkpoint(i, &batch),
152                    }
153                    .fail();
154                } else if true_count > 1 {
155                    return CheckpointOverlappedSnafu {
156                        checkpoint: self.remap_checkpoint(i, &batch),
157                    }
158                    .fail();
159                }
160            }
161        }
162
163        Ok(())
164    }
165
166    /// Remap the normalized checkpoint data to the original values.
167    fn remap_checkpoint(&self, i: usize, batch: &RecordBatch) -> String {
168        let normalized_row = batch
169            .columns()
170            .iter()
171            .map(|col| {
172                let array = col.as_any().downcast_ref::<Float64Array>().unwrap();
173                array.value(i)
174            })
175            .collect::<Vec<_>>();
176
177        let mut check_point = String::new();
178        let schema = batch.schema();
179        for (col_index, normalized_value) in normalized_row.iter().enumerate() {
180            let dim = schema
181                .field(col_index)
182                .name()
183                .parse::<usize>()
184                .ok()
185                .and_then(|index| self.dims.get(index));
186            let col_name = dim.map(ToString::to_string).unwrap_or_default();
187
188            if col_index > 0 {
189                check_point.push_str(", ");
190            }
191
192            // Check if point is on NORMALIZE_STEP or between steps
193            if let Some(values) = dim.and_then(|dim| self.collider.normalized_values.get(dim)) {
194                let normalize_step = NORMALIZE_STEP.0;
195
196                // Check if the normalized value is on a NORMALIZE_STEP boundary
197                let remainder = normalized_value % normalize_step;
198                let is_on_step = remainder.abs() < f64::EPSILON
199                    || (normalize_step - remainder).abs() < f64::EPSILON * 2.0;
200
201                if is_on_step {
202                    let index = (normalized_value / normalize_step).round() as usize;
203                    if index < values.len() {
204                        let original_value = &values[index].0;
205                        check_point.push_str(&format!("{}={}", col_name, original_value));
206                    } else {
207                        check_point.push_str(&format!("{}=unknown", col_name));
208                    }
209                } else {
210                    let lower_index = (normalized_value / normalize_step).floor() as usize;
211                    let upper_index = (normalized_value / normalize_step).ceil() as usize;
212
213                    // Handle edge cases: value is outside the valid range
214                    if lower_index == upper_index && lower_index == 0 {
215                        // Value is less than the first value
216                        let first_original = &values[0].0;
217                        check_point.push_str(&format!("{}<{}", col_name, first_original));
218                    } else if upper_index == values.len() {
219                        // Value is greater than the last value
220                        let last_original = &values[values.len() - 1].0;
221                        check_point.push_str(&format!("{}>{}", col_name, last_original));
222                    } else {
223                        // Normal case: value is between two valid values
224                        let lower_original = if lower_index < values.len() {
225                            values[lower_index].0.to_string()
226                        } else {
227                            "unknown".to_string()
228                        };
229
230                        let upper_original = if upper_index < values.len() {
231                            values[upper_index].0.to_string()
232                        } else {
233                            "unknown".to_string()
234                        };
235
236                        check_point.push_str(&format!(
237                            "{}<{}<{}",
238                            lower_original, col_name, upper_original
239                        ));
240                    }
241                }
242            } else {
243                // Fallback if column not found in normalized values
244                check_point.push_str(&format!("{}:unknown", col_name));
245            }
246        }
247
248        check_point
249    }
250}
251
252/// Generates a point matrix that contains permutations of `matrix_foundation`'s values
253struct MatrixGenerator {
254    matrix_foundation: HashMap<String, Vec<OrderedF64>>,
255    // Iterator state
256    current_index: usize,
257    schema: Schema,
258    column_names: Vec<String>,
259    // Preprocessed attributes
260    /// Total number of combinations of `matrix_foundation`'s values
261    total_combinations: usize,
262    /// Biased suffix product of `matrix_foundation`'s values
263    ///
264    /// The i-th element is the product of the sizes of all columns after the i-th column.
265    /// For example, if `matrix_foundation` is `{"a": [1, 2, 3], "b": [4, 5, 6]}`,
266    /// then `biased_suffix_product` is `[3, 1]`.
267    biased_suffix_product: Vec<usize>,
268}
269
270const MAX_BATCH_SIZE: usize = 8192;
271
272impl MatrixGenerator {
273    pub fn new(matrix_foundation: HashMap<&str, Vec<OrderedF64>>) -> Self {
274        // Convert to owned HashMap to avoid lifetime issues
275        let owned_matrix_foundation: HashMap<String, Vec<OrderedF64>> = matrix_foundation
276            .into_iter()
277            .map(|(k, v)| (k.to_string(), v))
278            .collect();
279
280        let mut fields = owned_matrix_foundation
281            .keys()
282            .map(|k| Field::new(k.clone(), DataType::Float64, false))
283            .collect::<Vec<_>>();
284        fields.sort_unstable();
285        let schema = Schema::new(fields.clone());
286
287        // Store column names in the same order as fields
288        let column_names: Vec<String> = fields.iter().map(|field| field.name().clone()).collect();
289
290        // Calculate total number of combinations and suffix product
291        let mut biased_suffix_product = Vec::with_capacity(column_names.len() + 1);
292        let mut product = 1;
293        biased_suffix_product.push(product);
294        for col_name in column_names.iter().rev() {
295            product *= owned_matrix_foundation[col_name].len();
296            biased_suffix_product.push(product);
297        }
298        biased_suffix_product.pop();
299        biased_suffix_product.reverse();
300
301        Self {
302            matrix_foundation: owned_matrix_foundation,
303            current_index: 0,
304            schema,
305            column_names,
306            total_combinations: product,
307            biased_suffix_product,
308        }
309    }
310
311    pub fn schema(&self) -> &Schema {
312        &self.schema
313    }
314
315    fn generate_batch(&self, start_index: usize, batch_size: usize) -> RecordBatch {
316        let actual_batch_size = batch_size.min(self.total_combinations - start_index);
317
318        // Create array builders
319        let mut array_builders: Vec<Float64Builder> = Vec::with_capacity(self.column_names.len());
320        for _ in 0..self.column_names.len() {
321            array_builders.push(Float64Builder::with_capacity(actual_batch_size));
322        }
323
324        // Generate combinations for this batch
325        for combination_offset in 0..actual_batch_size {
326            let combination_index = start_index + combination_offset;
327
328            // For each column, determine which value to use for this combination
329            for (col_idx, col_name) in self.column_names.iter().enumerate() {
330                let values = &self.matrix_foundation[col_name];
331                let stride = self.biased_suffix_product[col_idx];
332                let value_index = (combination_index / stride) % values.len();
333                let value = *values[value_index].as_ref();
334
335                array_builders[col_idx].append_value(value);
336            }
337        }
338
339        // Finish arrays and create RecordBatch
340        let arrays: Vec<_> = array_builders
341            .into_iter()
342            .map(|mut builder| Arc::new(builder.finish()) as _)
343            .collect();
344
345        RecordBatch::try_new(Arc::new(self.schema.clone()), arrays)
346            .expect("Failed to create RecordBatch from generated arrays")
347    }
348}
349
350impl Iterator for MatrixGenerator {
351    type Item = RecordBatch;
352
353    fn next(&mut self) -> Option<Self::Item> {
354        if self.current_index >= self.total_combinations {
355            return None;
356        }
357
358        let remaining = self.total_combinations - self.current_index;
359        let batch_size = remaining.min(MAX_BATCH_SIZE);
360
361        let batch = self.generate_batch(self.current_index, batch_size);
362        self.current_index += batch_size;
363
364        Some(batch)
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use std::collections::HashMap;
371
372    use datatypes::value::Value;
373
374    use super::*;
375    use crate::expr::col;
376    use crate::multi_dim::MultiDimPartitionRule;
377
378    #[test]
379    fn test_hash_and_same_named_column_are_distinct() {
380        use crate::expr::Operand;
381        use crate::function::IcebergHash;
382
383        let hash = IcebergHash::new("host");
384        // A column can be named like the hash call, but it is another dimension, so these
385        // two rules don't cover each other's complement.
386        let result = MultiDimPartitionRule::try_new(
387            vec!["host".to_string(), hash.to_string()],
388            vec![1, 2],
389            vec![
390                Operand::IcebergHash(hash.clone()).lt(Value::Int32(0)),
391                col(hash.to_string()).gt_eq(Value::Int32(0)),
392            ],
393            true,
394        );
395        assert!(result.is_err());
396    }
397
398    #[test]
399    fn test_matrix_generator_single_column() {
400        let mut matrix_foundation = HashMap::new();
401        matrix_foundation.insert(
402            "col1",
403            vec![
404                OrderedF64::from(1.0),
405                OrderedF64::from(2.0),
406                OrderedF64::from(3.0),
407            ],
408        );
409
410        let mut generator = MatrixGenerator::new(matrix_foundation);
411        let batch = generator.next().unwrap();
412
413        assert_eq!(batch.num_rows(), 3);
414        assert_eq!(batch.num_columns(), 1);
415        assert_eq!(batch.schema().field(0).name(), "col1");
416
417        let col1_array = batch
418            .column(0)
419            .as_any()
420            .downcast_ref::<datatypes::arrow::array::Float64Array>()
421            .unwrap();
422        assert_eq!(col1_array.value(0), 1.0);
423        assert_eq!(col1_array.value(1), 2.0);
424        assert_eq!(col1_array.value(2), 3.0);
425
426        // Should be no more batches for such a small dataset
427        assert!(generator.next().is_none());
428    }
429
430    #[test]
431    fn test_matrix_generator_three_columns_cartesian_product() {
432        let mut matrix_foundation = HashMap::new();
433        matrix_foundation.insert("a", vec![OrderedF64::from(1.0), OrderedF64::from(2.0)]);
434        matrix_foundation.insert("b", vec![OrderedF64::from(10.0), OrderedF64::from(20.0)]);
435        matrix_foundation.insert(
436            "c",
437            vec![
438                OrderedF64::from(100.0),
439                OrderedF64::from(200.0),
440                OrderedF64::from(300.0),
441            ],
442        );
443
444        let mut generator = MatrixGenerator::new(matrix_foundation);
445        let batch = generator.next().unwrap();
446
447        // Should have 2 * 2 * 3 = 12 combinations
448        assert_eq!(batch.num_rows(), 12);
449        assert_eq!(batch.num_columns(), 3);
450
451        let a_array = batch
452            .column(0)
453            .as_any()
454            .downcast_ref::<datatypes::arrow::array::Float64Array>()
455            .unwrap();
456        let b_array = batch
457            .column(1)
458            .as_any()
459            .downcast_ref::<datatypes::arrow::array::Float64Array>()
460            .unwrap();
461        let c_array = batch
462            .column(2)
463            .as_any()
464            .downcast_ref::<datatypes::arrow::array::Float64Array>()
465            .unwrap();
466
467        // Verify first few combinations (a changes slowest, c changes fastest)
468        let expected = vec![
469            (1.0, 10.0, 100.0),
470            (1.0, 10.0, 200.0),
471            (1.0, 10.0, 300.0),
472            (1.0, 20.0, 100.0),
473            (1.0, 20.0, 200.0),
474            (1.0, 20.0, 300.0),
475            (2.0, 10.0, 100.0),
476            (2.0, 10.0, 200.0),
477            (2.0, 10.0, 300.0),
478            (2.0, 20.0, 100.0),
479            (2.0, 20.0, 200.0),
480            (2.0, 20.0, 300.0),
481        ];
482        #[allow(clippy::needless_range_loop)]
483        for i in 0..batch.num_rows() {
484            assert_eq!(
485                (a_array.value(i), b_array.value(i), c_array.value(i)),
486                expected[i]
487            );
488        }
489
490        // Should be no more batches for such a small dataset
491        assert!(generator.next().is_none());
492    }
493
494    #[test]
495    fn test_matrix_generator_iterator_small_batches() {
496        let mut matrix_foundation = HashMap::new();
497        matrix_foundation.insert("col1", vec![OrderedF64::from(1.0), OrderedF64::from(2.0)]);
498        matrix_foundation.insert(
499            "col2",
500            vec![
501                OrderedF64::from(10.0),
502                OrderedF64::from(20.0),
503                OrderedF64::from(30.0),
504            ],
505        );
506
507        let generator = MatrixGenerator::new(matrix_foundation);
508
509        // Total combinations should be 2 * 3 = 6
510        assert_eq!(generator.total_combinations, 6);
511
512        let mut total_rows = 0;
513
514        for batch in generator {
515            total_rows += batch.num_rows();
516            assert_eq!(batch.num_columns(), 2);
517
518            // Verify each batch is valid
519            assert!(batch.num_rows() > 0);
520            assert!(batch.num_rows() <= MAX_BATCH_SIZE);
521        }
522
523        assert_eq!(total_rows, 6);
524    }
525
526    #[test]
527    fn test_matrix_generator_empty_column_values() {
528        let mut matrix_foundation = HashMap::new();
529        matrix_foundation.insert("col1", vec![]);
530
531        let mut generator = MatrixGenerator::new(matrix_foundation);
532
533        // Should have 0 total combinations when any column is empty
534        assert_eq!(generator.total_combinations, 0);
535
536        // Should have no batches when total combinations is 0
537        assert!(generator.next().is_none());
538    }
539
540    #[test]
541    fn test_matrix_generator_large_dataset_batching() {
542        // Create a dataset that will exceed MAX_BATCH_SIZE (8192)
543        // 20 * 20 * 21 = 8400 > 8192
544        let mut matrix_foundation = HashMap::new();
545
546        let values1: Vec<OrderedF64> = (0..20).map(|i| OrderedF64::from(i as f64)).collect();
547        let values2: Vec<OrderedF64> = (0..20)
548            .map(|i| OrderedF64::from(i as f64 + 100.0))
549            .collect();
550        let values3: Vec<OrderedF64> = (0..21)
551            .map(|i| OrderedF64::from(i as f64 + 1000.0))
552            .collect();
553
554        matrix_foundation.insert("col1", values1);
555        matrix_foundation.insert("col2", values2);
556        matrix_foundation.insert("col3", values3);
557
558        let generator = MatrixGenerator::new(matrix_foundation);
559
560        assert_eq!(generator.total_combinations, 8400);
561
562        let mut total_rows = 0;
563        let mut batch_count = 0;
564        let mut first_batch_size = None;
565
566        for batch in generator {
567            batch_count += 1;
568            let batch_size = batch.num_rows();
569            total_rows += batch_size;
570
571            if first_batch_size.is_none() {
572                first_batch_size = Some(batch_size);
573            }
574
575            // Each batch should not exceed MAX_BATCH_SIZE
576            assert!(batch_size <= MAX_BATCH_SIZE);
577            assert_eq!(batch.num_columns(), 3);
578        }
579
580        assert_eq!(total_rows, 8400);
581        assert!(batch_count > 1);
582        assert_eq!(first_batch_size.unwrap(), MAX_BATCH_SIZE);
583    }
584
585    #[test]
586    fn test_remap_checkpoint_values() {
587        // Create rule with single column
588        let rule = MultiDimPartitionRule::try_new(
589            vec!["host".to_string(), "value".to_string()],
590            vec![1, 2, 3],
591            vec![
592                col("host")
593                    .lt(Value::Int64(0))
594                    .and(col("value").lt(Value::Int64(0))),
595                col("host")
596                    .lt(Value::Int64(0))
597                    .and(col("value").gt_eq(Value::Int64(0))),
598                col("host")
599                    .gt_eq(Value::Int64(0))
600                    .and(col("host").lt(Value::Int64(1)))
601                    .and(col("value").lt(Value::Int64(1))),
602                col("host")
603                    .gt_eq(Value::Int64(0))
604                    .and(col("host").lt(Value::Int64(1)))
605                    .and(col("value").gt_eq(Value::Int64(1))),
606                col("host")
607                    .gt_eq(Value::Int64(1))
608                    .and(col("host").lt(Value::Int64(2)))
609                    .and(col("value").lt(Value::Int64(2))),
610                col("host")
611                    .gt_eq(Value::Int64(1))
612                    .and(col("host").lt(Value::Int64(2)))
613                    .and(col("value").gt_eq(Value::Int64(2))),
614                col("host")
615                    .gt_eq(Value::Int64(2))
616                    .and(col("host").lt(Value::Int64(3)))
617                    .and(col("value").lt(Value::Int64(3))),
618                col("host")
619                    .gt_eq(Value::Int64(2))
620                    .and(col("host").lt(Value::Int64(3)))
621                    .and(col("value").gt_eq(Value::Int64(3))),
622                col("host").gt_eq(Value::Int64(3)),
623            ],
624            true,
625        )
626        .unwrap();
627        let checker = PartitionChecker::try_new(&rule).unwrap();
628
629        let schema = Arc::new(Schema::new(vec![
630            Field::new("0", DataType::Float64, false),
631            Field::new("1", DataType::Float64, false),
632        ]));
633        let host_array = Float64Array::from(vec![-0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]);
634        let value_array = Float64Array::from(vec![-0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]);
635        let batch = RecordBatch::try_new(schema, vec![Arc::new(host_array), Arc::new(value_array)])
636            .unwrap();
637
638        let checkpoint = checker.remap_checkpoint(0, &batch);
639        assert_eq!(checkpoint, "host<0, value<0");
640        let checkpoint = checker.remap_checkpoint(1, &batch);
641        assert_eq!(checkpoint, "host=0, value=0");
642        let checkpoint = checker.remap_checkpoint(6, &batch);
643        assert_eq!(checkpoint, "2<host<3, 2<value<3");
644        let checkpoint = checker.remap_checkpoint(7, &batch);
645        assert_eq!(checkpoint, "host=3, value=3");
646        let checkpoint = checker.remap_checkpoint(8, &batch);
647        assert_eq!(checkpoint, "host>3, value>3");
648    }
649}
partition/checker.rs

partition/
checker.rs
