partition/
manager.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::sync::Arc;
17
18use api::v1::Rows;
19use common_meta::cache::{TableRoute, TableRouteCacheRef};
20use common_meta::key::table_route::{PhysicalTableRouteValue, TableRouteManager};
21use common_meta::kv_backend::KvBackendRef;
22use common_meta::peer::Peer;
23use common_meta::rpc::router::{self, RegionRoute};
24use snafu::{OptionExt, ResultExt};
25use store_api::storage::{RegionId, RegionNumber};
26use table::metadata::{TableId, TableInfo};
27
28use crate::cache::{CachedPartitionInfo, PartitionInfoCacheRef, PhysicalPartitionInfo};
29use crate::error::{FindLeaderSnafu, Result};
30use crate::expr::PartitionExpr;
31use crate::multi_dim::MultiDimPartitionRule;
32use crate::splitter::RowSplitter;
33use crate::{PartitionRuleRef, error};
34
35pub type PartitionRuleManagerRef = Arc<PartitionRuleManager>;
36
37/// PartitionRuleManager manages the table routes and partition rules.
38/// It provides methods to find regions by:
39/// - values (in case of insertion)
40/// - filters (in case of select, deletion and update)
41pub struct PartitionRuleManager {
42    table_route_manager: TableRouteManager,
43    table_route_cache: TableRouteCacheRef,
44    partition_info_cache: PartitionInfoCacheRef,
45}
46
47#[derive(Debug, Clone)]
48pub struct PartitionInfo {
49    pub id: RegionId,
50    pub partition_expr: Option<PartitionExpr>,
51}
52
53#[derive(Debug, Clone)]
54pub struct PartitionInfoWithVersion {
55    pub id: RegionId,
56    pub partition_expr: Option<PartitionExpr>,
57    pub partition_expr_version: Option<u64>,
58}
59
60impl PartitionRuleManager {
61    pub fn new(
62        kv_backend: KvBackendRef,
63        table_route_cache: TableRouteCacheRef,
64        partition_info_cache: PartitionInfoCacheRef,
65    ) -> Self {
66        Self {
67            table_route_manager: TableRouteManager::new(kv_backend),
68            table_route_cache,
69            partition_info_cache,
70        }
71    }
72
73    pub async fn find_physical_table_route(
74        &self,
75        table_id: TableId,
76    ) -> Result<Arc<PhysicalTableRouteValue>> {
77        Ok(self.find_physical_table_route_with_id(table_id).await?.1)
78    }
79
80    /// Returns the resolved physical table id and its route.
81    pub async fn find_physical_table_route_with_id(
82        &self,
83        table_id: TableId,
84    ) -> Result<(TableId, Arc<PhysicalTableRouteValue>)> {
85        match self
86            .table_route_cache
87            .get(table_id)
88            .await
89            .context(error::TableRouteManagerSnafu)?
90            .context(error::TableRouteNotFoundSnafu { table_id })?
91            .as_ref()
92        {
93            TableRoute::Physical(physical_table_route) => {
94                Ok((table_id, physical_table_route.clone()))
95            }
96            TableRoute::Logical(logical_table_route) => {
97                let physical_table_id = logical_table_route.physical_table_id();
98                let physical_table_route = self
99                    .table_route_cache
100                    .get(physical_table_id)
101                    .await
102                    .context(error::TableRouteManagerSnafu)?
103                    .context(error::TableRouteNotFoundSnafu {
104                        table_id: physical_table_id,
105                    })?;
106
107                let physical_table_route = physical_table_route
108                    .as_physical_table_route_ref()
109                    .context(error::UnexpectedSnafu{
110                        err_msg: format!(
111                            "Expected the physical table route, but got logical table route, table: {physical_table_id}"
112                        ),
113                    })?;
114
115                Ok((physical_table_id, physical_table_route.clone()))
116            }
117        }
118    }
119
120    pub async fn batch_find_region_routes(
121        &self,
122        table_ids: &[TableId],
123    ) -> Result<HashMap<TableId, Vec<RegionRoute>>> {
124        let table_routes = self
125            .table_route_manager
126            .batch_get_physical_table_routes(table_ids)
127            .await
128            .context(error::TableRouteManagerSnafu)?;
129
130        let mut table_region_routes = HashMap::with_capacity(table_routes.len());
131
132        for (table_id, table_route) in table_routes {
133            let region_routes = table_route.region_routes;
134            table_region_routes.insert(table_id, region_routes);
135        }
136
137        Ok(table_region_routes)
138    }
139
140    /// Returns the physical partition info with version.
141    pub async fn find_physical_partition_info(
142        &self,
143        table_id: TableId,
144    ) -> Result<Arc<PhysicalPartitionInfo>> {
145        let cached = self
146            .partition_info_cache
147            .get(table_id)
148            .await
149            .context(error::GetPartitionInfoSnafu)?
150            .context(error::TableRouteNotFoundSnafu { table_id })?;
151        match cached {
152            CachedPartitionInfo::Physical(info) => Ok(info),
153            CachedPartitionInfo::Logical(physical_table_id) => {
154                let cached = self
155                    .partition_info_cache
156                    .get(physical_table_id)
157                    .await
158                    .context(error::GetPartitionInfoSnafu)?
159                    .context(error::TableRouteNotFoundSnafu {
160                        table_id: physical_table_id,
161                    })?;
162                let info = cached.into_physical().context(error::UnexpectedSnafu{
163                        err_msg: format!(
164                            "Expected the physical partition info, but got logical partable route, table: {physical_table_id}"
165                        )
166                    })?;
167
168                Ok(info)
169            }
170        }
171    }
172
173    pub async fn batch_find_table_partitions(
174        &self,
175        table_ids: &[TableId],
176    ) -> Result<HashMap<TableId, Vec<PartitionInfo>>> {
177        let batch_region_routes = self.batch_find_region_routes(table_ids).await?;
178
179        let mut results = HashMap::with_capacity(table_ids.len());
180
181        for (table_id, region_routes) in batch_region_routes {
182            results.insert(
183                table_id,
184                create_partitions_from_region_routes(table_id, &region_routes)?,
185            );
186        }
187
188        Ok(results)
189    }
190
191    pub async fn find_table_partition_rule(
192        &self,
193        table_info: &TableInfo,
194    ) -> Result<(PartitionRuleRef, HashMap<RegionNumber, Option<u64>>)> {
195        let partition_columns = table_info
196            .meta
197            .partition_column_names()
198            .cloned()
199            .collect::<Vec<_>>();
200
201        let partition_info = self
202            .find_physical_partition_info(table_info.table_id())
203            .await?;
204        let partition_versions = partition_info
205            .partitions
206            .iter()
207            .map(|r| (r.id.region_number(), r.partition_expr_version))
208            .collect::<HashMap<RegionNumber, Option<u64>>>();
209        let regions = partition_info
210            .partitions
211            .iter()
212            .map(|x| x.id.region_number())
213            .collect::<Vec<RegionNumber>>();
214        let exprs = partition_info
215            .partitions
216            .iter()
217            .filter_map(|x| x.partition_expr.as_ref())
218            .cloned()
219            .collect::<Vec<_>>();
220        let partition_rule = Arc::new(MultiDimPartitionRule::try_new(
221            partition_columns,
222            regions,
223            exprs,
224            false,
225        )?) as _;
226        Ok((partition_rule, partition_versions))
227    }
228
229    /// Find the leader of the region.
230    pub async fn find_region_leader(&self, region_id: RegionId) -> Result<Peer> {
231        let region_routes = &self
232            .find_physical_table_route(region_id.table_id())
233            .await?
234            .region_routes;
235
236        router::find_region_leader(region_routes, region_id.region_number()).context(
237            FindLeaderSnafu {
238                region_id,
239                table_id: region_id.table_id(),
240            },
241        )
242    }
243
244    pub async fn split_rows(
245        &self,
246        table_info: &TableInfo,
247        rows: Rows,
248    ) -> Result<HashMap<RegionNumber, (Rows, Option<u64>)>> {
249        let (partition_rule, partition_versions) =
250            self.find_table_partition_rule(table_info).await?;
251
252        let result = RowSplitter::new(partition_rule)
253            .split(rows)?
254            .into_iter()
255            .map(|(region_number, rows)| {
256                (
257                    region_number,
258                    (
259                        rows,
260                        partition_versions
261                            .get(&region_number)
262                            .copied()
263                            .unwrap_or_default(),
264                    ),
265                )
266            })
267            .collect::<HashMap<_, _>>();
268
269        Ok(result)
270    }
271}
272
273/// Creates partitions from region routes.
274pub fn create_partitions_from_region_routes(
275    table_id: TableId,
276    region_routes: &[RegionRoute],
277) -> Result<Vec<PartitionInfo>> {
278    let mut partitions = Vec::with_capacity(region_routes.len());
279    for r in region_routes {
280        let partition_expr = PartitionExpr::from_json_str(&r.region.partition_expr())?;
281
282        // The region routes belong to the physical table but are shared among all logical tables.
283        // That it to say, the region id points to the physical table, so we need to use the actual
284        // table id (which may be a logical table) to renew the region id.
285        let id = RegionId::new(table_id, r.region.id.region_number());
286        partitions.push(PartitionInfo { id, partition_expr });
287    }
288
289    Ok(partitions)
290}
partition/manager.rs

partition/
manager.rs
