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

partition/
multi_dim.rs
