partition/
multi_dim.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::any::Any;
16use std::cmp::Ordering;
17use std::collections::HashMap;
18use std::collections::hash_map::Entry;
19use std::sync::{Arc, RwLock};
20
21use datafusion_expr::ColumnarValue;
22use datafusion_physical_expr::PhysicalExpr;
23use datatypes::arrow;
24use datatypes::arrow::array::{Array, BooleanArray, BooleanBufferBuilder, RecordBatch};
25use datatypes::arrow::buffer::BooleanBuffer;
26use datatypes::arrow::compute::prep_null_mask_filter;
27use datatypes::arrow::datatypes::Schema;
28use datatypes::prelude::Value;
29use datatypes::vectors::{Helper, VectorRef};
30use serde::{Deserialize, Serialize};
31use snafu::{OptionExt, ResultExt, ensure};
32use store_api::storage::RegionNumber;
33
34use crate::PartitionRule;
35use crate::checker::PartitionChecker;
36use crate::error::{self, Result, UndefinedColumnSnafu};
37use crate::expr::{Operand, PartitionExpr, RestrictedOp};
38use crate::partition::RegionMask;
39
40/// The default region number when no partition exprs are matched.
41const DEFAULT_REGION: RegionNumber = 0;
42
43type PhysicalExprCache = Option<(Vec<Arc<dyn PhysicalExpr>>, Arc<Schema>)>;
44
45/// Multi-Dimiension partition rule. RFC [here](https://github.com/GreptimeTeam/greptimedb/blob/main/docs/rfcs/2024-02-21-multi-dimension-partition-rule/rfc.md)
46///
47/// This partition rule is defined by a set of simple expressions on the partition
48/// key columns. Compare to RANGE partition, which can be considered as
49/// single-dimension rule, this will evaluate expression on each column separately.
50#[derive(Debug, Serialize, Deserialize)]
51pub struct MultiDimPartitionRule {
52    /// Allow list of which columns can be used for partitioning.
53    partition_columns: Vec<String>,
54    /// Name to index of `partition_columns`. Used for quick lookup.
55    name_to_index: HashMap<String, usize>,
56    /// Region number for each partition. This list has the same length as `exprs`
57    /// (dispiting the default region).
58    regions: Vec<RegionNumber>,
59    /// Partition expressions.
60    exprs: Vec<PartitionExpr>,
61    /// Cache of physical expressions.
62    #[serde(skip)]
63    physical_expr_cache: RwLock<PhysicalExprCache>,
64}
65
66impl MultiDimPartitionRule {
67    /// Create a new [`MultiDimPartitionRule`].
68    ///
69    /// If `check_exprs` is true, the function will check if the expressions are valid. This is
70    /// required when constructing a new partition rule like `CREATE TABLE` or `ALTER TABLE`.
71    pub fn try_new(
72        partition_columns: Vec<String>,
73        regions: Vec<RegionNumber>,
74        exprs: Vec<PartitionExpr>,
75        check_exprs: bool,
76    ) -> Result<Self> {
77        let name_to_index = partition_columns
78            .iter()
79            .enumerate()
80            .map(|(i, name)| (name.clone(), i))
81            .collect::<HashMap<_, _>>();
82
83        let rule = Self {
84            partition_columns,
85            name_to_index,
86            regions,
87            exprs,
88            physical_expr_cache: RwLock::new(None),
89        };
90
91        if check_exprs {
92            let checker = PartitionChecker::try_new(&rule)?;
93            checker.check()?;
94        }
95
96        Ok(rule)
97    }
98
99    pub fn exprs(&self) -> &[PartitionExpr] {
100        &self.exprs
101    }
102
103    fn find_region(&self, values: &[Value]) -> Result<RegionNumber> {
104        ensure!(
105            values.len() == self.partition_columns.len(),
106            error::RegionKeysSizeSnafu {
107                expect: self.partition_columns.len(),
108                actual: values.len(),
109            }
110        );
111
112        for (region_index, expr) in self.exprs.iter().enumerate() {
113            if self.evaluate_expr(expr, values)? {
114                return Ok(self.regions[region_index]);
115            }
116        }
117
118        // return the default region number
119        Ok(DEFAULT_REGION)
120    }
121
122    fn evaluate_expr(&self, expr: &PartitionExpr, values: &[Value]) -> Result<bool> {
123        match (expr.lhs.as_ref(), expr.rhs.as_ref()) {
124            (Operand::Column(name), Operand::Value(r)) => {
125                let index = self
126                    .name_to_index
127                    .get(name)
128                    .context(UndefinedColumnSnafu { column: name })?;
129                let l = &values[*index];
130                Self::perform_op(l, &expr.op, r)
131            }
132            (Operand::Value(l), Operand::Column(name)) => {
133                let index = self
134                    .name_to_index
135                    .get(name)
136                    .context(UndefinedColumnSnafu { column: name })?;
137                let r = &values[*index];
138                Self::perform_op(l, &expr.op, r)
139            }
140            (Operand::Expr(lhs), Operand::Expr(rhs)) => {
141                let lhs = self.evaluate_expr(lhs, values)?;
142                let rhs = self.evaluate_expr(rhs, values)?;
143                match expr.op {
144                    RestrictedOp::And => Ok(lhs && rhs),
145                    RestrictedOp::Or => Ok(lhs || rhs),
146                    _ => unreachable!(),
147                }
148            }
149            _ => unreachable!(),
150        }
151    }
152
153    fn perform_op(lhs: &Value, op: &RestrictedOp, rhs: &Value) -> Result<bool> {
154        let result = match op {
155            RestrictedOp::Eq => lhs.eq(rhs),
156            RestrictedOp::NotEq => lhs.ne(rhs),
157            RestrictedOp::Lt => lhs.partial_cmp(rhs) == Some(Ordering::Less),
158            RestrictedOp::LtEq => {
159                let result = lhs.partial_cmp(rhs);
160                result == Some(Ordering::Less) || result == Some(Ordering::Equal)
161            }
162            RestrictedOp::Gt => lhs.partial_cmp(rhs) == Some(Ordering::Greater),
163            RestrictedOp::GtEq => {
164                let result = lhs.partial_cmp(rhs);
165                result == Some(Ordering::Greater) || result == Some(Ordering::Equal)
166            }
167            RestrictedOp::And | RestrictedOp::Or => unreachable!(),
168        };
169
170        Ok(result)
171    }
172
173    pub fn row_at(&self, cols: &[VectorRef], index: usize, row: &mut [Value]) -> Result<()> {
174        for (col_idx, col) in cols.iter().enumerate() {
175            row[col_idx] = col.get(index);
176        }
177        Ok(())
178    }
179
180    pub fn record_batch_to_cols(&self, record_batch: &RecordBatch) -> Result<Vec<VectorRef>> {
181        self.partition_columns
182            .iter()
183            .map(|col_name| {
184                record_batch
185                    .column_by_name(col_name)
186                    .context(UndefinedColumnSnafu { column: col_name })
187                    .and_then(|array| {
188                        Helper::try_into_vector(array).context(error::ConvertToVectorSnafu)
189                    })
190            })
191            .collect::<Result<Vec<_>>>()
192    }
193
194    pub fn split_record_batch_naive(
195        &self,
196        record_batch: &RecordBatch,
197    ) -> Result<HashMap<RegionNumber, BooleanArray>> {
198        let num_rows = record_batch.num_rows();
199
200        let mut result = self
201            .regions
202            .iter()
203            .map(|region| {
204                let mut builder = BooleanBufferBuilder::new(num_rows);
205                builder.append_n(num_rows, false);
206                (*region, builder)
207            })
208            .collect::<HashMap<_, _>>();
209
210        let cols = self.record_batch_to_cols(record_batch)?;
211        let mut current_row = vec![Value::Null; self.partition_columns.len()];
212        for row_idx in 0..num_rows {
213            self.row_at(&cols, row_idx, &mut current_row)?;
214            let current_region = self.find_region(&current_row)?;
215            let region_mask = result
216                .get_mut(&current_region)
217                .unwrap_or_else(|| panic!("Region {} must be initialized", current_region));
218            region_mask.set_bit(row_idx, true);
219        }
220
221        Ok(result
222            .into_iter()
223            .map(|(region, mut mask)| (region, BooleanArray::new(mask.finish(), None)))
224            .collect())
225    }
226
227    pub fn split_record_batch(
228        &self,
229        record_batch: &RecordBatch,
230    ) -> Result<HashMap<RegionNumber, RegionMask>> {
231        let num_rows = record_batch.num_rows();
232        if self.regions.len() == 1 {
233            return Ok([(
234                self.regions[0],
235                RegionMask::from(BooleanArray::from(vec![true; num_rows])),
236            )]
237            .into_iter()
238            .collect());
239        }
240        let physical_exprs = {
241            let cache_read_guard = self.physical_expr_cache.read().unwrap();
242            if let Some((cached_exprs, schema)) = cache_read_guard.as_ref()
243                && schema == record_batch.schema_ref()
244            {
245                cached_exprs.clone()
246            } else {
247                drop(cache_read_guard); // Release the read lock before acquiring write lock
248
249                let schema = record_batch.schema();
250                let new_cache = self
251                    .exprs
252                    .iter()
253                    .map(|e| e.try_as_physical_expr(&schema))
254                    .collect::<Result<Vec<_>>>()?;
255
256                let mut cache_write_guard = self.physical_expr_cache.write().unwrap();
257                cache_write_guard.replace((new_cache.clone(), schema));
258                new_cache
259            }
260        };
261
262        let mut result: HashMap<u32, RegionMask> = physical_exprs
263            .iter()
264            .zip(self.regions.iter())
265            .filter_map(|(expr, region_num)| {
266                let col_val = match expr
267                    .evaluate(record_batch)
268                    .context(error::EvaluateRecordBatchSnafu)
269                {
270                    Ok(array) => array,
271                    Err(e) => {
272                        return Some(Err(e));
273                    }
274                };
275                let array = match columnar_value_to_boolean_array(col_val, num_rows) {
276                    // A NULL result means the row is not in this region. Masks are merged
277                    // with non-Kleene `or` below, where a NULL would drop the row.
278                    Ok(array) if array.null_count() > 0 => prep_null_mask_filter(&array),
279                    Ok(array) => array,
280                    Err(e) => {
281                        return Some(Err(e));
282                    }
283                };
284                let selected_rows = array.true_count();
285                if selected_rows == 0 {
286                    // skip empty region in results.
287                    return None;
288                }
289                Some(Ok((*region_num, RegionMask::new(array, selected_rows))))
290            })
291            .collect::<error::Result<_>>()?;
292
293        let selected = if result.len() == 1 {
294            result.values().next().unwrap().array().clone()
295        } else {
296            let mut selected = BooleanArray::new(BooleanBuffer::new_unset(num_rows), None);
297            for region_mask in result.values() {
298                selected = arrow::compute::kernels::boolean::or(&selected, region_mask.array())
299                    .context(error::ComputeArrowKernelSnafu)?;
300            }
301            selected
302        };
303
304        // fast path: all rows are selected
305        if selected.true_count() == num_rows {
306            return Ok(result);
307        }
308
309        // find unselected rows and assign to default region
310        let unselected = arrow::compute::kernels::boolean::not(&selected)
311            .context(error::ComputeArrowKernelSnafu)?;
312        match result.entry(DEFAULT_REGION) {
313            Entry::Occupied(mut o) => {
314                // merge default region with unselected rows.
315                let default_region_mask = RegionMask::from(
316                    arrow::compute::kernels::boolean::or(o.get().array(), &unselected)
317                        .context(error::ComputeArrowKernelSnafu)?,
318                );
319                o.insert(default_region_mask);
320            }
321            Entry::Vacant(v) => {
322                // default region has no rows, simply put all unselected rows to default region.
323                v.insert(RegionMask::from(unselected));
324            }
325        }
326        Ok(result)
327    }
328}
329
330fn columnar_value_to_boolean_array(
331    col_val: ColumnarValue,
332    num_rows: usize,
333) -> Result<BooleanArray> {
334    let column = col_val
335        .into_array(num_rows)
336        .context(error::EvaluateRecordBatchSnafu)?;
337    let array = column
338        .as_any()
339        .downcast_ref::<BooleanArray>()
340        .with_context(|| error::UnexpectedColumnTypeSnafu {
341            data_type: column.data_type().clone(),
342        })?;
343    Ok(array.clone())
344}
345
346impl PartitionRule for MultiDimPartitionRule {
347    fn as_any(&self) -> &dyn Any {
348        self
349    }
350
351    fn partition_columns(&self) -> &[String] {
352        &self.partition_columns
353    }
354
355    fn find_region(&self, values: &[Value]) -> Result<RegionNumber> {
356        self.find_region(values)
357    }
358
359    fn split_record_batch(
360        &self,
361        record_batch: &RecordBatch,
362    ) -> Result<HashMap<RegionNumber, RegionMask>> {
363        self.split_record_batch(record_batch)
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    use std::assert_matches;
370
371    use super::*;
372    use crate::error::{self, Error};
373    use crate::expr::col;
374
375    #[test]
376    fn test_find_region() {
377        // PARTITION ON COLUMNS (b) (
378        //     b < 'hz',
379        //     b >= 'hz' AND b < 'sh',
380        //     b >= 'sh'
381        // )
382        let rule = MultiDimPartitionRule::try_new(
383            vec!["b".to_string()],
384            vec![1, 2, 3],
385            vec![
386                PartitionExpr::new(
387                    Operand::Column("b".to_string()),
388                    RestrictedOp::Lt,
389                    Operand::Value(datatypes::value::Value::String("hz".into())),
390                ),
391                PartitionExpr::new(
392                    Operand::Expr(PartitionExpr::new(
393                        Operand::Column("b".to_string()),
394                        RestrictedOp::GtEq,
395                        Operand::Value(datatypes::value::Value::String("hz".into())),
396                    )),
397                    RestrictedOp::And,
398                    Operand::Expr(PartitionExpr::new(
399                        Operand::Column("b".to_string()),
400                        RestrictedOp::Lt,
401                        Operand::Value(datatypes::value::Value::String("sh".into())),
402                    )),
403                ),
404                PartitionExpr::new(
405                    Operand::Column("b".to_string()),
406                    RestrictedOp::GtEq,
407                    Operand::Value(datatypes::value::Value::String("sh".into())),
408                ),
409            ],
410            true,
411        )
412        .unwrap();
413        assert_matches!(
414            rule.find_region(&["foo".into(), 1000_i32.into()]),
415            Err(error::Error::RegionKeysSize {
416                expect: 1,
417                actual: 2,
418                ..
419            })
420        );
421        assert_matches!(rule.find_region(&["foo".into()]), Ok(1));
422        assert_matches!(rule.find_region(&["bar".into()]), Ok(1));
423        assert_matches!(rule.find_region(&["hz".into()]), Ok(2));
424        assert_matches!(rule.find_region(&["hzz".into()]), Ok(2));
425        assert_matches!(rule.find_region(&["sh".into()]), Ok(3));
426        assert_matches!(rule.find_region(&["zzzz".into()]), Ok(3));
427    }
428
429    #[test]
430    fn test_find_region_rejects_undeclared_column_on_lhs() {
431        let rule = MultiDimPartitionRule::try_new(
432            vec!["host".to_string()],
433            vec![1],
434            vec![PartitionExpr::new(
435                Operand::Column("rack".to_string()),
436                RestrictedOp::Lt,
437                Operand::Value(Value::String("n".into())),
438            )],
439            false,
440        )
441        .unwrap();
442
443        assert_matches!(
444            rule.find_region(&[Value::String("z".into())]),
445            Err(Error::UndefinedColumn { column, .. }) if column == "rack"
446        );
447    }
448
449    #[test]
450    fn test_find_region_rejects_undeclared_column_on_rhs() {
451        let rule = MultiDimPartitionRule::try_new(
452            vec!["host".to_string()],
453            vec![1],
454            vec![PartitionExpr::new(
455                Operand::Value(Value::String("n".into())),
456                RestrictedOp::Gt,
457                Operand::Column("rack".to_string()),
458            )],
459            false,
460        )
461        .unwrap();
462
463        assert_matches!(
464            rule.find_region(&[Value::String("z".into())]),
465            Err(Error::UndefinedColumn { column, .. }) if column == "rack"
466        );
467    }
468
469    #[test]
470    fn invalid_expr_case_1() {
471        // PARTITION ON COLUMNS (b) (
472        //     b <= b >= 'hz' AND b < 'sh',
473        // )
474        let rule = MultiDimPartitionRule::try_new(
475            vec!["a".to_string(), "b".to_string()],
476            vec![1],
477            vec![PartitionExpr::new(
478                Operand::Column("b".to_string()),
479                RestrictedOp::LtEq,
480                Operand::Expr(PartitionExpr::new(
481                    Operand::Expr(PartitionExpr::new(
482                        Operand::Column("b".to_string()),
483                        RestrictedOp::GtEq,
484                        Operand::Value(datatypes::value::Value::String("hz".into())),
485                    )),
486                    RestrictedOp::And,
487                    Operand::Expr(PartitionExpr::new(
488                        Operand::Column("b".to_string()),
489                        RestrictedOp::Lt,
490                        Operand::Value(datatypes::value::Value::String("sh".into())),
491                    )),
492                )),
493            )],
494            true,
495        );
496
497        // check rule
498        assert_matches!(rule.unwrap_err(), Error::InvalidExpr { .. });
499    }
500
501    #[test]
502    fn invalid_expr_case_2() {
503        // PARTITION ON COLUMNS (b) (
504        //     b >= 'hz' AND 'sh',
505        // )
506        let rule = MultiDimPartitionRule::try_new(
507            vec!["a".to_string(), "b".to_string()],
508            vec![1],
509            vec![PartitionExpr::new(
510                Operand::Expr(PartitionExpr::new(
511                    Operand::Column("b".to_string()),
512                    RestrictedOp::GtEq,
513                    Operand::Value(datatypes::value::Value::String("hz".into())),
514                )),
515                RestrictedOp::And,
516                Operand::Value(datatypes::value::Value::String("sh".into())),
517            )],
518            true,
519        );
520
521        // check rule
522        assert_matches!(rule.unwrap_err(), Error::InvalidExpr { .. });
523    }
524
525    /// ```ignore
526    ///          │          │
527    ///          │          │
528    /// ─────────┼──────────┼────────────► b
529    ///          │          │
530    ///          │          │
531    ///      b <= h     b >= s
532    /// ```
533    #[test]
534    fn empty_expr_case_1() {
535        // PARTITION ON COLUMNS (b) (
536        //     b <= 'h',
537        //     b >= 's'
538        // )
539        let rule = MultiDimPartitionRule::try_new(
540            vec!["a".to_string(), "b".to_string()],
541            vec![1, 2],
542            vec![
543                PartitionExpr::new(
544                    Operand::Column("b".to_string()),
545                    RestrictedOp::LtEq,
546                    Operand::Value(datatypes::value::Value::String("h".into())),
547                ),
548                PartitionExpr::new(
549                    Operand::Column("b".to_string()),
550                    RestrictedOp::GtEq,
551                    Operand::Value(datatypes::value::Value::String("s".into())),
552                ),
553            ],
554            true,
555        );
556
557        // check rule
558        assert_matches!(rule.unwrap_err(), Error::CheckpointNotCovered { .. });
559    }
560
561    /// ```
562    ///     a
563    ///     ▲
564    ///     │                   ‖
565    ///     │                   ‖
566    /// 200 │         ┌─────────┤
567    ///     │         │         │
568    ///     │         │         │
569    ///     │         │         │
570    /// 100 │   ======┴─────────┘
571    ///     │
572    ///     └──────────────────────────►b
573    ///              10          20
574    /// ```
575    #[test]
576    fn empty_expr_case_2() {
577        // PARTITION ON COLUMNS (b) (
578        //     a >= 100 AND b <= 10  OR  a > 100 AND a <= 200 AND b <= 10  OR  a >= 200 AND b > 10 AND b <= 20  OR  a > 200 AND b <= 20
579        //     a < 100 AND b <= 20  OR  a >= 100 AND b > 20
580        // )
581        let rule = MultiDimPartitionRule::try_new(
582            vec!["a".to_string(), "b".to_string()],
583            vec![1, 2],
584            vec![
585                PartitionExpr::new(
586                    Operand::Expr(PartitionExpr::new(
587                        Operand::Expr(PartitionExpr::new(
588                            //  a >= 100 AND b <= 10
589                            Operand::Expr(PartitionExpr::new(
590                                Operand::Expr(PartitionExpr::new(
591                                    Operand::Column("a".to_string()),
592                                    RestrictedOp::GtEq,
593                                    Operand::Value(datatypes::value::Value::Int64(100)),
594                                )),
595                                RestrictedOp::And,
596                                Operand::Expr(PartitionExpr::new(
597                                    Operand::Column("b".to_string()),
598                                    RestrictedOp::LtEq,
599                                    Operand::Value(datatypes::value::Value::Int64(10)),
600                                )),
601                            )),
602                            RestrictedOp::Or,
603                            // a > 100 AND a <= 200 AND b <= 10
604                            Operand::Expr(PartitionExpr::new(
605                                Operand::Expr(PartitionExpr::new(
606                                    Operand::Expr(PartitionExpr::new(
607                                        Operand::Column("a".to_string()),
608                                        RestrictedOp::Gt,
609                                        Operand::Value(datatypes::value::Value::Int64(100)),
610                                    )),
611                                    RestrictedOp::And,
612                                    Operand::Expr(PartitionExpr::new(
613                                        Operand::Column("a".to_string()),
614                                        RestrictedOp::LtEq,
615                                        Operand::Value(datatypes::value::Value::Int64(200)),
616                                    )),
617                                )),
618                                RestrictedOp::And,
619                                Operand::Expr(PartitionExpr::new(
620                                    Operand::Column("b".to_string()),
621                                    RestrictedOp::LtEq,
622                                    Operand::Value(datatypes::value::Value::Int64(10)),
623                                )),
624                            )),
625                        )),
626                        RestrictedOp::Or,
627                        // a >= 200 AND b > 10 AND b <= 20
628                        Operand::Expr(PartitionExpr::new(
629                            Operand::Expr(PartitionExpr::new(
630                                Operand::Expr(PartitionExpr::new(
631                                    Operand::Column("a".to_string()),
632                                    RestrictedOp::GtEq,
633                                    Operand::Value(datatypes::value::Value::Int64(200)),
634                                )),
635                                RestrictedOp::And,
636                                Operand::Expr(PartitionExpr::new(
637                                    Operand::Column("b".to_string()),
638                                    RestrictedOp::Gt,
639                                    Operand::Value(datatypes::value::Value::Int64(10)),
640                                )),
641                            )),
642                            RestrictedOp::And,
643                            Operand::Expr(PartitionExpr::new(
644                                Operand::Column("b".to_string()),
645                                RestrictedOp::LtEq,
646                                Operand::Value(datatypes::value::Value::Int64(20)),
647                            )),
648                        )),
649                    )),
650                    RestrictedOp::Or,
651                    // a > 200 AND b <= 20
652                    Operand::Expr(PartitionExpr::new(
653                        Operand::Expr(PartitionExpr::new(
654                            Operand::Column("a".to_string()),
655                            RestrictedOp::Gt,
656                            Operand::Value(datatypes::value::Value::Int64(200)),
657                        )),
658                        RestrictedOp::And,
659                        Operand::Expr(PartitionExpr::new(
660                            Operand::Column("b".to_string()),
661                            RestrictedOp::LtEq,
662                            Operand::Value(datatypes::value::Value::Int64(20)),
663                        )),
664                    )),
665                ),
666                PartitionExpr::new(
667                    // a < 100 AND b <= 20
668                    Operand::Expr(PartitionExpr::new(
669                        Operand::Expr(PartitionExpr::new(
670                            Operand::Column("a".to_string()),
671                            RestrictedOp::Lt,
672                            Operand::Value(datatypes::value::Value::Int64(100)),
673                        )),
674                        RestrictedOp::And,
675                        Operand::Expr(PartitionExpr::new(
676                            Operand::Column("b".to_string()),
677                            RestrictedOp::LtEq,
678                            Operand::Value(datatypes::value::Value::Int64(20)),
679                        )),
680                    )),
681                    RestrictedOp::Or,
682                    // a >= 100 AND b > 20
683                    Operand::Expr(PartitionExpr::new(
684                        Operand::Expr(PartitionExpr::new(
685                            Operand::Column("a".to_string()),
686                            RestrictedOp::GtEq,
687                            Operand::Value(datatypes::value::Value::Int64(100)),
688                        )),
689                        RestrictedOp::And,
690                        Operand::Expr(PartitionExpr::new(
691                            Operand::Column("b".to_string()),
692                            RestrictedOp::GtEq,
693                            Operand::Value(datatypes::value::Value::Int64(20)),
694                        )),
695                    )),
696                ),
697            ],
698            true,
699        );
700
701        // check rule
702        assert_matches!(rule.unwrap_err(), Error::CheckpointNotCovered { .. });
703    }
704
705    #[test]
706    fn duplicate_expr_case_1() {
707        // PARTITION ON COLUMNS (a) (
708        //     a <= 20,
709        //     a >= 10
710        // )
711        let rule = MultiDimPartitionRule::try_new(
712            vec!["a".to_string(), "b".to_string()],
713            vec![1, 2],
714            vec![
715                PartitionExpr::new(
716                    Operand::Column("a".to_string()),
717                    RestrictedOp::LtEq,
718                    Operand::Value(datatypes::value::Value::Int64(20)),
719                ),
720                PartitionExpr::new(
721                    Operand::Column("a".to_string()),
722                    RestrictedOp::GtEq,
723                    Operand::Value(datatypes::value::Value::Int64(10)),
724                ),
725            ],
726            true,
727        );
728
729        // check rule
730        assert_matches!(rule.unwrap_err(), Error::CheckpointOverlapped { .. });
731    }
732
733    #[test]
734    fn duplicate_expr_case_2() {
735        // PARTITION ON COLUMNS (a) (
736        //     a != 20,
737        //     a <= 20,
738        //     a > 20,
739        // )
740        let rule = MultiDimPartitionRule::try_new(
741            vec!["a".to_string(), "b".to_string()],
742            vec![1, 2],
743            vec![
744                PartitionExpr::new(
745                    Operand::Column("a".to_string()),
746                    RestrictedOp::NotEq,
747                    Operand::Value(datatypes::value::Value::Int64(20)),
748                ),
749                PartitionExpr::new(
750                    Operand::Column("a".to_string()),
751                    RestrictedOp::LtEq,
752                    Operand::Value(datatypes::value::Value::Int64(20)),
753                ),
754                PartitionExpr::new(
755                    Operand::Column("a".to_string()),
756                    RestrictedOp::Gt,
757                    Operand::Value(datatypes::value::Value::Int64(20)),
758                ),
759            ],
760            true,
761        );
762
763        // check rule
764        assert_matches!(rule.unwrap_err(), Error::CheckpointOverlapped { .. });
765    }
766
767    /// ```ignore
768    /// value
769    ///                                 │
770    ///                                 │
771    ///    value=10 --------------------│
772    ///                                 │
773    /// ────────────────────────────────┼──► host
774    ///                                 │
775    ///                             host=server10
776    /// ```
777    #[test]
778    fn test_partial_divided() {
779        let _rule = MultiDimPartitionRule::try_new(
780            vec!["host".to_string(), "value".to_string()],
781            vec![0, 1, 2, 3],
782            vec![
783                col("host")
784                    .lt(Value::String("server10".into()))
785                    .and(col("value").lt(Value::Int64(10))),
786                col("host")
787                    .lt(Value::String("server10".into()))
788                    .and(col("value").gt_eq(Value::Int64(10))),
789                col("host").gt_eq(Value::String("server10".into())),
790            ],
791            true,
792        )
793        .unwrap();
794    }
795}
796
797#[cfg(test)]
798mod test_split_record_batch {
799    use std::sync::Arc;
800
801    use common_time::timestamp::Timestamp;
802    use datafusion_common::ScalarValue;
803    use datatypes::arrow::array::{Int64Array, StringArray, TimestampNanosecondArray};
804    use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
805    use datatypes::arrow::record_batch::RecordBatch;
806    use rand::Rng;
807
808    use super::*;
809    use crate::expr::{Operand, col};
810
811    fn test_schema() -> Arc<Schema> {
812        Arc::new(Schema::new(vec![
813            Field::new("host", DataType::Utf8, false),
814            Field::new("value", DataType::Int64, false),
815        ]))
816    }
817
818    fn generate_random_record_batch(num_rows: usize) -> RecordBatch {
819        let schema = test_schema();
820        let mut rng = rand::thread_rng();
821        let mut host_array = Vec::with_capacity(num_rows);
822        let mut value_array = Vec::with_capacity(num_rows);
823        for _ in 0..num_rows {
824            host_array.push(format!("server{}", rng.gen_range(0..20)));
825            value_array.push(rng.gen_range(0..20));
826        }
827        let host_array = StringArray::from(host_array);
828        let value_array = Int64Array::from(value_array);
829        RecordBatch::try_new(schema, vec![Arc::new(host_array), Arc::new(value_array)]).unwrap()
830    }
831
832    #[test]
833    fn test_split_record_batch_by_one_column() {
834        // Create a simple MultiDimPartitionRule
835        let rule = MultiDimPartitionRule::try_new(
836            vec!["host".to_string(), "value".to_string()],
837            vec![0, 1],
838            vec![
839                col("host").lt(Value::String("server1".into())),
840                col("host").gt_eq(Value::String("server1".into())),
841            ],
842            true,
843        )
844        .unwrap();
845
846        let batch = generate_random_record_batch(1000);
847        // Split the batch
848        let result = rule.split_record_batch(&batch).unwrap();
849        let expected = rule.split_record_batch_naive(&batch).unwrap();
850        assert_eq!(result.len(), expected.len());
851        for (region, value) in &result {
852            assert_eq!(
853                value.array(),
854                expected.get(region).unwrap(),
855                "failed on region: {}",
856                region
857            );
858        }
859    }
860
861    #[test]
862    fn test_split_record_batch_widened_time_index_unit() {
863        // https://github.com/GreptimeTeam/greptimedb/issues/9747
864        // The partition bounds keep the unit they were created with (millisecond)
865        // after the time index column is widened to nanosecond, so the batch
866        // schema and the bound values mix timestamp units.
867        let bound = Value::from(Timestamp::new_millisecond(1_000));
868        let rule = MultiDimPartitionRule::try_new(
869            vec!["ts".to_string()],
870            vec![0, 1],
871            vec![col("ts").lt(bound.clone()), col("ts").gt_eq(bound)],
872            true,
873        )
874        .unwrap();
875
876        let schema = Arc::new(Schema::new(vec![Field::new(
877            "ts",
878            DataType::Timestamp(TimeUnit::Nanosecond, None),
879            false,
880        )]));
881        let batch = RecordBatch::try_new(
882            schema,
883            vec![Arc::new(TimestampNanosecondArray::from(vec![
884                500_000_000,
885                2_000_000_000,
886            ]))],
887        )
888        .unwrap();
889
890        // 500_000_000ns < 1000ms and 2_000_000_000ns >= 1000ms: one row per region.
891        let expected = rule.split_record_batch_naive(&batch).unwrap();
892        assert_eq!(expected.len(), 2);
893        let result = rule.split_record_batch(&batch).unwrap();
894        assert_eq!(result.len(), expected.len());
895        for (region, value) in &result {
896            assert_eq!(value.array(), expected.get(region).unwrap());
897        }
898    }
899
900    #[test]
901    fn test_split_record_batch_empty() {
902        // Create a simple MultiDimPartitionRule
903        let rule = MultiDimPartitionRule::try_new(
904            vec!["host".to_string()],
905            vec![1],
906            vec![
907                col("host").lt(Value::String("server1".into())),
908                col("host").gt_eq(Value::String("server1".into())),
909            ],
910            true,
911        )
912        .unwrap();
913
914        let schema = test_schema();
915        let host_array = StringArray::from(Vec::<&str>::new());
916        let value_array = Int64Array::from(Vec::<i64>::new());
917        let batch = RecordBatch::try_new(schema, vec![Arc::new(host_array), Arc::new(value_array)])
918            .unwrap();
919
920        let result = rule.split_record_batch(&batch).unwrap();
921        assert_eq!(result.len(), 1);
922    }
923
924    #[test]
925    fn test_split_record_batch_by_two_columns() {
926        let rule = MultiDimPartitionRule::try_new(
927            vec!["host".to_string(), "value".to_string()],
928            vec![0, 1, 2, 3],
929            vec![
930                col("host")
931                    .lt(Value::String("server10".into()))
932                    .and(col("value").lt(Value::Int64(10))),
933                col("host")
934                    .lt(Value::String("server10".into()))
935                    .and(col("value").gt_eq(Value::Int64(10))),
936                col("host")
937                    .gt_eq(Value::String("server10".into()))
938                    .and(col("value").lt(Value::Int64(10))),
939                col("host")
940                    .gt_eq(Value::String("server10".into()))
941                    .and(col("value").gt_eq(Value::Int64(10))),
942            ],
943            true,
944        )
945        .unwrap();
946
947        let batch = generate_random_record_batch(1000);
948        let result = rule.split_record_batch(&batch).unwrap();
949        let expected = rule.split_record_batch_naive(&batch).unwrap();
950        assert_eq!(result.len(), expected.len());
951        for (region, value) in &result {
952            assert_eq!(value.array(), expected.get(region).unwrap());
953        }
954    }
955
956    #[test]
957    fn test_all_rows_selected() {
958        // Test the fast path where all rows are selected by some partition
959        let rule = MultiDimPartitionRule::try_new(
960            vec!["value".to_string()],
961            vec![1, 2],
962            vec![
963                col("value").lt(Value::Int64(30)),
964                col("value").gt_eq(Value::Int64(30)),
965            ],
966            true,
967        )
968        .unwrap();
969
970        let schema = test_schema();
971        let host_array = StringArray::from(vec!["server1", "server2", "server3", "server4"]);
972        let value_array = Int64Array::from(vec![10, 20, 30, 40]);
973        let batch = RecordBatch::try_new(schema, vec![Arc::new(host_array), Arc::new(value_array)])
974            .unwrap();
975
976        let result = rule.split_record_batch(&batch).unwrap();
977
978        // Check that we have 2 regions and no default region
979        assert_eq!(result.len(), 2);
980        assert!(result.contains_key(&1));
981        assert!(result.contains_key(&2));
982
983        // Verify each region has the correct number of rows
984        assert_eq!(result.get(&1).unwrap().selected_rows(), 2); // values < 30
985        assert_eq!(result.get(&2).unwrap().selected_rows(), 2); // values >= 30
986    }
987
988    #[test]
989    fn test_split_record_batch_with_scalar_predicate() {
990        // Ensure split handles conjunctive/disjunctive predicates on the same column.
991        let rule = MultiDimPartitionRule::try_new(
992            vec!["host".to_string()],
993            vec![0, 1],
994            vec![
995                PartitionExpr::new(
996                    Operand::Column("host".to_string()),
997                    RestrictedOp::Lt,
998                    Operand::Value(Value::String("never_happen_1".into())),
999                ),
1000                PartitionExpr::new(
1001                    Operand::Expr(PartitionExpr::new(
1002                        Operand::Column("host".to_string()),
1003                        RestrictedOp::GtEq,
1004                        Operand::Value(Value::String("never_happen_1".into())),
1005                    )),
1006                    RestrictedOp::And,
1007                    Operand::Value(Value::Boolean(false)),
1008                ),
1009            ],
1010            false,
1011        )
1012        .unwrap();
1013
1014        let batch = generate_random_record_batch(8);
1015        let result = rule.split_record_batch(&batch).unwrap();
1016
1017        assert_eq!(result.len(), 1);
1018        assert!(result.contains_key(&0));
1019
1020        let total_rows = result.get(&0).unwrap().selected_rows();
1021        assert_eq!(total_rows, batch.num_rows());
1022    }
1023
1024    #[test]
1025    fn test_comparison_batch_matches_row() {
1026        use datatypes::arrow::array::Array;
1027        use datatypes::arrow::datatypes::{DataType, Field};
1028
1029        let inputs = [None, Some(4), Some(5), Some(6)];
1030        let batch = RecordBatch::try_new(
1031            Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])),
1032            vec![Arc::new(Int64Array::from(inputs.to_vec()))],
1033        )
1034        .unwrap();
1035        let mut mismatches = vec![];
1036        for op in [
1037            RestrictedOp::Eq,
1038            RestrictedOp::NotEq,
1039            RestrictedOp::Lt,
1040            RestrictedOp::LtEq,
1041            RestrictedOp::Gt,
1042            RestrictedOp::GtEq,
1043        ] {
1044            for bound in [Value::Null, Value::Int64(5)] {
1045                // Both operand orders, since a rule restored without canonicalization may
1046                // keep the literal on the left.
1047                for expr in [
1048                    PartitionExpr::new(col("a"), op.clone(), bound.clone().into()),
1049                    PartitionExpr {
1050                        lhs: Box::new(bound.clone().into()),
1051                        op: op.clone(),
1052                        rhs: Box::new(col("a")),
1053                    },
1054                ] {
1055                    let rule = MultiDimPartitionRule::try_new(
1056                        vec!["a".to_string()],
1057                        vec![1],
1058                        vec![expr.clone()],
1059                        false,
1060                    )
1061                    .unwrap();
1062                    let by_batch = match expr
1063                        .try_as_physical_expr(batch.schema_ref())
1064                        .unwrap()
1065                        .evaluate(&batch)
1066                    {
1067                        Ok(by_batch) => {
1068                            columnar_value_to_boolean_array(by_batch, batch.num_rows()).unwrap()
1069                        }
1070                        Err(err) => {
1071                            mismatches.push(format!("{expr}: batch error {err}"));
1072                            continue;
1073                        }
1074                    };
1075                    for (row, input) in inputs.iter().enumerate() {
1076                        let value = input.map_or(Value::Null, Value::Int64);
1077                        let by_row = rule.evaluate_expr(&expr, &[value]).unwrap();
1078                        let by_batch = by_batch.is_valid(row) && by_batch.value(row);
1079                        if by_row != by_batch {
1080                            mismatches.push(format!(
1081                                "{expr} with a = {input:?}: row {by_row}, batch {by_batch}"
1082                            ));
1083                        }
1084                    }
1085                }
1086            }
1087        }
1088        assert!(mismatches.is_empty(), "{mismatches:#?}");
1089    }
1090
1091    #[test]
1092    fn test_split_null_rows_match_row_routing() {
1093        use datatypes::arrow::datatypes::{DataType, Field};
1094
1095        let rule = MultiDimPartitionRule::try_new(
1096            vec!["host".to_string()],
1097            vec![0, 1],
1098            vec![
1099                col("host").not_eq(Value::from("a")),
1100                col("host").eq(Value::from("a")),
1101            ],
1102            true,
1103        )
1104        .unwrap();
1105        for hosts in [vec![Some("a"), None, Some("z")], vec![Some("a"), None]] {
1106            let batch = RecordBatch::try_new(
1107                Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])),
1108                vec![Arc::new(StringArray::from(hosts.clone()))],
1109            )
1110            .unwrap();
1111            let by_row = rule.split_record_batch_naive(&batch).unwrap();
1112            let by_batch = rule.split_record_batch(&batch).unwrap();
1113            for (region, expected) in by_row {
1114                let actual = by_batch
1115                    .get(&region)
1116                    .map(|mask| mask.array().clone())
1117                    .unwrap_or_else(|| BooleanArray::from(vec![false; batch.num_rows()]));
1118                assert_eq!(actual, expected, "{hosts:?}, region {region}");
1119            }
1120        }
1121    }
1122
1123    #[test]
1124    fn test_columnar_value_to_boolean_array_scalar_false() {
1125        let result = columnar_value_to_boolean_array(
1126            ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))),
1127            4,
1128        )
1129        .unwrap();
1130        assert_eq!(result.len(), 4);
1131        assert_eq!(result.true_count(), 0);
1132    }
1133
1134    #[test]
1135    fn test_columnar_value_to_boolean_array_scalar_true() {
1136        let result = columnar_value_to_boolean_array(
1137            ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))),
1138            4,
1139        )
1140        .unwrap();
1141        assert_eq!(result.len(), 4);
1142        assert_eq!(result.true_count(), 4);
1143    }
1144}
partition/multi_dim.rs

partition/
multi_dim.rs
