partition/
partition.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::collections::HashMap;
17use std::fmt::Debug;
18use std::sync::Arc;
19
20use datatypes::arrow::array::{BooleanArray, RecordBatch};
21use datatypes::prelude::Value;
22use serde::{Deserialize, Serialize};
23use store_api::storage::RegionNumber;
24
25use crate::error::Result;
26
27pub type PartitionRuleRef = Arc<dyn PartitionRule>;
28
29pub trait PartitionRule: Sync + Send {
30    fn as_any(&self) -> &dyn Any;
31
32    fn partition_columns(&self) -> Vec<String>;
33
34    /// Finds the target region by the partition values.
35    ///
36    /// Note that the `values` should have the same length as the `partition_columns`.
37    fn find_region(&self, values: &[Value]) -> Result<RegionNumber>;
38
39    /// Split the record batch into multiple regions by the partition values.
40    /// The result is a map from region mask in which the array is true for the rows that match the partition values.
41    /// Region with now rows selected may not appear in result map.
42    fn split_record_batch(
43        &self,
44        record_batch: &RecordBatch,
45    ) -> Result<HashMap<RegionNumber, RegionMask>>;
46}
47
48/// The bound of one partition.
49#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
50pub enum PartitionBound {
51    /// Deprecated since 0.9.0.
52    Value(Value),
53    /// Deprecated since 0.15.0.
54    MaxValue,
55    Expr(crate::expr::PartitionExpr),
56}
57
58pub struct RegionMask {
59    array: BooleanArray,
60    selected_rows: usize,
61}
62
63impl From<BooleanArray> for RegionMask {
64    fn from(array: BooleanArray) -> Self {
65        let selected_rows = array.true_count();
66        Self {
67            array,
68            selected_rows,
69        }
70    }
71}
72
73impl RegionMask {
74    pub fn new(array: BooleanArray, selected_rows: usize) -> Self {
75        Self {
76            array,
77            selected_rows,
78        }
79    }
80
81    pub fn array(&self) -> &BooleanArray {
82        &self.array
83    }
84
85    /// All rows are selected.
86    pub fn select_all(&self) -> bool {
87        self.selected_rows == self.array.len()
88    }
89
90    /// No row is selected.
91    pub fn select_none(&self) -> bool {
92        self.selected_rows == 0
93    }
94
95    pub fn selected_rows(&self) -> usize {
96        self.selected_rows
97    }
98}
partition/partition.rs

partition/
partition.rs
