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

partition/
multi_dim.rs
