mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-01-10 15:22:56 +00:00
459 lines
16 KiB
Rust
459 lines
16 KiB
Rust
// Copyright 2023 Greptime Team
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
use std::sync::Arc;
|
|
|
|
use common_error::ext::BoxedError;
|
|
use common_recordbatch::SendableRecordBatchStream;
|
|
use store_api::data_source::DataSource;
|
|
use store_api::storage::ScanRequest;
|
|
use table::metadata::{FilterPushDownType, TableInfoRef};
|
|
use table::thin_table::{ThinTable, ThinTableAdapter};
|
|
use table::TableRef;
|
|
|
|
use crate::error::NotSupportedSnafu;
|
|
|
|
pub mod delete;
|
|
pub mod insert;
|
|
|
|
#[derive(Clone)]
|
|
pub struct DistTable;
|
|
|
|
impl DistTable {
|
|
pub fn table(table_info: TableInfoRef) -> TableRef {
|
|
let thin_table = ThinTable::new(table_info, FilterPushDownType::Inexact);
|
|
let data_source = Arc::new(DummyDataSource);
|
|
Arc::new(ThinTableAdapter::new(thin_table, data_source))
|
|
}
|
|
}
|
|
|
|
pub struct DummyDataSource;
|
|
|
|
impl DataSource for DummyDataSource {
|
|
fn get_stream(&self, _request: ScanRequest) -> Result<SendableRecordBatchStream, BoxedError> {
|
|
NotSupportedSnafu {
|
|
feat: "get stream from a distributed table",
|
|
}
|
|
.fail()
|
|
.map_err(BoxedError::new)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) mod test {
|
|
use std::collections::BTreeMap;
|
|
use std::sync::atomic::{AtomicU32, Ordering};
|
|
|
|
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
|
|
use common_meta::peer::Peer;
|
|
use common_meta::rpc::router::{Region, RegionRoute, Table, TableRoute};
|
|
use common_meta::table_name::TableName;
|
|
use common_query::prelude::Expr;
|
|
use datafusion_expr::expr_fn::{and, binary_expr, col, or};
|
|
use datafusion_expr::{lit, Operator};
|
|
use meta_client::client::MetaClient;
|
|
use meter_core::collect::Collect;
|
|
use meter_core::data::{ReadRecord, WriteRecord};
|
|
use meter_core::global::global_registry;
|
|
use meter_core::write_calc::WriteCalculator;
|
|
use partition::columns::RangeColumnsPartitionRule;
|
|
use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
|
|
use partition::partition::{PartitionBound, PartitionDef};
|
|
use partition::range::RangePartitionRule;
|
|
use partition::route::TableRoutes;
|
|
use partition::PartitionRuleRef;
|
|
use store_api::storage::RegionNumber;
|
|
use table::meter_insert_request;
|
|
use table::requests::InsertRequest;
|
|
|
|
use super::*;
|
|
|
|
/// Create a partition rule manager with two tables, one is partitioned by single column, and
|
|
/// the other one is two. The tables are under default catalog and schema.
|
|
///
|
|
/// Table named "one_column_partitioning_table" is partitioned by column "a" like this:
|
|
/// PARTITION BY RANGE (a) (
|
|
/// PARTITION r1 VALUES LESS THAN (10),
|
|
/// PARTITION r2 VALUES LESS THAN (50),
|
|
/// PARTITION r3 VALUES LESS THAN (MAXVALUE),
|
|
/// )
|
|
///
|
|
/// Table named "two_column_partitioning_table" is partitioned by columns "a" and "b" like this:
|
|
/// PARTITION BY RANGE (a, b) (
|
|
/// PARTITION r1 VALUES LESS THAN (10, 'hz'),
|
|
/// PARTITION r2 VALUES LESS THAN (50, 'sh'),
|
|
/// PARTITION r3 VALUES LESS THAN (MAXVALUE, MAXVALUE),
|
|
/// )
|
|
pub(crate) async fn create_partition_rule_manager() -> PartitionRuleManagerRef {
|
|
let table_routes = Arc::new(TableRoutes::new(Arc::new(MetaClient::default())));
|
|
let partition_manager = Arc::new(PartitionRuleManager::new(table_routes.clone()));
|
|
|
|
let table_name = TableName::new(
|
|
DEFAULT_CATALOG_NAME,
|
|
DEFAULT_SCHEMA_NAME,
|
|
"one_column_partitioning_table",
|
|
);
|
|
let table_route = TableRoute::new(
|
|
Table {
|
|
id: 1,
|
|
table_name: table_name.clone(),
|
|
table_schema: vec![],
|
|
},
|
|
vec![
|
|
RegionRoute {
|
|
region: Region {
|
|
id: 3.into(),
|
|
name: "r1".to_string(),
|
|
partition: Some(
|
|
PartitionDef::new(
|
|
vec!["a".to_string()],
|
|
vec![PartitionBound::Value(10_i32.into())],
|
|
)
|
|
.try_into()
|
|
.unwrap(),
|
|
),
|
|
attrs: BTreeMap::new(),
|
|
},
|
|
leader_peer: Some(Peer::new(3, "")),
|
|
follower_peers: vec![],
|
|
},
|
|
RegionRoute {
|
|
region: Region {
|
|
id: 2.into(),
|
|
name: "r2".to_string(),
|
|
partition: Some(
|
|
PartitionDef::new(
|
|
vec!["a".to_string()],
|
|
vec![PartitionBound::Value(50_i32.into())],
|
|
)
|
|
.try_into()
|
|
.unwrap(),
|
|
),
|
|
attrs: BTreeMap::new(),
|
|
},
|
|
leader_peer: Some(Peer::new(2, "")),
|
|
follower_peers: vec![],
|
|
},
|
|
RegionRoute {
|
|
region: Region {
|
|
id: 1.into(),
|
|
name: "r3".to_string(),
|
|
partition: Some(
|
|
PartitionDef::new(
|
|
vec!["a".to_string()],
|
|
vec![PartitionBound::MaxValue],
|
|
)
|
|
.try_into()
|
|
.unwrap(),
|
|
),
|
|
attrs: BTreeMap::new(),
|
|
},
|
|
leader_peer: Some(Peer::new(1, "")),
|
|
follower_peers: vec![],
|
|
},
|
|
],
|
|
);
|
|
table_routes
|
|
.insert_table_route(1, Arc::new(table_route))
|
|
.await;
|
|
|
|
let table_name = TableName::new(
|
|
DEFAULT_CATALOG_NAME,
|
|
DEFAULT_SCHEMA_NAME,
|
|
"two_column_partitioning_table",
|
|
);
|
|
let table_route = TableRoute::new(
|
|
Table {
|
|
id: 2,
|
|
table_name: table_name.clone(),
|
|
table_schema: vec![],
|
|
},
|
|
vec![
|
|
RegionRoute {
|
|
region: Region {
|
|
id: 1.into(),
|
|
name: "r1".to_string(),
|
|
partition: Some(
|
|
PartitionDef::new(
|
|
vec!["a".to_string(), "b".to_string()],
|
|
vec![
|
|
PartitionBound::Value(10_i32.into()),
|
|
PartitionBound::Value("hz".into()),
|
|
],
|
|
)
|
|
.try_into()
|
|
.unwrap(),
|
|
),
|
|
attrs: BTreeMap::new(),
|
|
},
|
|
leader_peer: None,
|
|
follower_peers: vec![],
|
|
},
|
|
RegionRoute {
|
|
region: Region {
|
|
id: 2.into(),
|
|
name: "r2".to_string(),
|
|
partition: Some(
|
|
PartitionDef::new(
|
|
vec!["a".to_string(), "b".to_string()],
|
|
vec![
|
|
PartitionBound::Value(50_i32.into()),
|
|
PartitionBound::Value("sh".into()),
|
|
],
|
|
)
|
|
.try_into()
|
|
.unwrap(),
|
|
),
|
|
attrs: BTreeMap::new(),
|
|
},
|
|
leader_peer: None,
|
|
follower_peers: vec![],
|
|
},
|
|
RegionRoute {
|
|
region: Region {
|
|
id: 3.into(),
|
|
name: "r3".to_string(),
|
|
partition: Some(
|
|
PartitionDef::new(
|
|
vec!["a".to_string(), "b".to_string()],
|
|
vec![PartitionBound::MaxValue, PartitionBound::MaxValue],
|
|
)
|
|
.try_into()
|
|
.unwrap(),
|
|
),
|
|
attrs: BTreeMap::new(),
|
|
},
|
|
leader_peer: None,
|
|
follower_peers: vec![],
|
|
},
|
|
],
|
|
);
|
|
table_routes
|
|
.insert_table_route(2, Arc::new(table_route))
|
|
.await;
|
|
|
|
partition_manager
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_find_partition_rule() {
|
|
let partition_manager = create_partition_rule_manager().await;
|
|
|
|
// "one_column_partitioning_table" has id 1
|
|
let partition_rule = partition_manager
|
|
.find_table_partition_rule(1)
|
|
.await
|
|
.unwrap();
|
|
let range_rule = partition_rule
|
|
.as_any()
|
|
.downcast_ref::<RangePartitionRule>()
|
|
.unwrap();
|
|
assert_eq!(range_rule.column_name(), "a");
|
|
assert_eq!(range_rule.all_regions(), &vec![3, 2, 1]);
|
|
assert_eq!(range_rule.bounds(), &vec![10_i32.into(), 50_i32.into()]);
|
|
|
|
// "two_column_partitioning_table" has table 2
|
|
let partition_rule = partition_manager
|
|
.find_table_partition_rule(2)
|
|
.await
|
|
.unwrap();
|
|
let range_columns_rule = partition_rule
|
|
.as_any()
|
|
.downcast_ref::<RangeColumnsPartitionRule>()
|
|
.unwrap();
|
|
assert_eq!(range_columns_rule.column_list(), &vec!["a", "b"]);
|
|
assert_eq!(
|
|
range_columns_rule.value_lists(),
|
|
&vec![
|
|
vec![
|
|
PartitionBound::Value(10_i32.into()),
|
|
PartitionBound::Value("hz".into()),
|
|
],
|
|
vec![
|
|
PartitionBound::Value(50_i32.into()),
|
|
PartitionBound::Value("sh".into()),
|
|
],
|
|
vec![PartitionBound::MaxValue, PartitionBound::MaxValue]
|
|
]
|
|
);
|
|
assert_eq!(range_columns_rule.regions(), &vec![1, 2, 3]);
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_find_regions() {
|
|
let partition_manager = Arc::new(PartitionRuleManager::new(Arc::new(TableRoutes::new(
|
|
Arc::new(MetaClient::default()),
|
|
))));
|
|
|
|
// PARTITION BY RANGE (a) (
|
|
// PARTITION r1 VALUES LESS THAN (10),
|
|
// PARTITION r2 VALUES LESS THAN (20),
|
|
// PARTITION r3 VALUES LESS THAN (50),
|
|
// PARTITION r4 VALUES LESS THAN (MAXVALUE),
|
|
// )
|
|
let partition_rule: PartitionRuleRef = Arc::new(RangePartitionRule::new(
|
|
"a",
|
|
vec![10_i32.into(), 20_i32.into(), 50_i32.into()],
|
|
vec![0_u32, 1, 2, 3],
|
|
)) as _;
|
|
|
|
let partition_rule_clone = partition_rule.clone();
|
|
let test = |filters: Vec<Expr>, expect_regions: Vec<RegionNumber>| {
|
|
let mut regions = partition_manager
|
|
.find_regions_by_filters(partition_rule_clone.clone(), filters.as_slice())
|
|
.unwrap();
|
|
regions.sort();
|
|
assert_eq!(regions, expect_regions);
|
|
};
|
|
|
|
// test simple filter
|
|
test(
|
|
vec![binary_expr(col("a"), Operator::Lt, lit(10)).into()], // a < 10
|
|
vec![0],
|
|
);
|
|
test(
|
|
vec![binary_expr(col("a"), Operator::LtEq, lit(10)).into()], // a <= 10
|
|
vec![0, 1],
|
|
);
|
|
test(
|
|
vec![binary_expr(lit(20), Operator::Gt, col("a")).into()], // 20 > a
|
|
vec![0, 1],
|
|
);
|
|
test(
|
|
vec![binary_expr(lit(20), Operator::GtEq, col("a")).into()], // 20 >= a
|
|
vec![0, 1, 2],
|
|
);
|
|
test(
|
|
vec![binary_expr(lit(45), Operator::Eq, col("a")).into()], // 45 == a
|
|
vec![2],
|
|
);
|
|
test(
|
|
vec![binary_expr(col("a"), Operator::NotEq, lit(45)).into()], // a != 45
|
|
vec![0, 1, 2, 3],
|
|
);
|
|
test(
|
|
vec![binary_expr(col("a"), Operator::Gt, lit(50)).into()], // a > 50
|
|
vec![3],
|
|
);
|
|
|
|
// test multiple filters
|
|
test(
|
|
vec![
|
|
binary_expr(col("a"), Operator::Gt, lit(10)).into(),
|
|
binary_expr(col("a"), Operator::Gt, lit(50)).into(),
|
|
], // [a > 10, a > 50]
|
|
vec![3],
|
|
);
|
|
|
|
// test finding all regions when provided with not supported filters or not partition column
|
|
test(
|
|
vec![binary_expr(col("row_id"), Operator::LtEq, lit(123)).into()], // row_id <= 123
|
|
vec![0, 1, 2, 3],
|
|
);
|
|
test(
|
|
vec![binary_expr(col("c"), Operator::Gt, lit(123)).into()], // c > 789
|
|
vec![0, 1, 2, 3],
|
|
);
|
|
|
|
// test complex "AND" or "OR" filters
|
|
test(
|
|
vec![and(
|
|
binary_expr(col("row_id"), Operator::Lt, lit(1)),
|
|
or(
|
|
binary_expr(col("row_id"), Operator::Lt, lit(1)),
|
|
binary_expr(col("a"), Operator::Lt, lit(1)),
|
|
),
|
|
)
|
|
.into()], // row_id < 1 OR (row_id < 1 AND a > 1)
|
|
vec![0, 1, 2, 3],
|
|
);
|
|
test(
|
|
vec![or(
|
|
binary_expr(col("a"), Operator::Lt, lit(20)),
|
|
binary_expr(col("a"), Operator::GtEq, lit(20)),
|
|
)
|
|
.into()], // a < 20 OR a >= 20
|
|
vec![0, 1, 2, 3],
|
|
);
|
|
test(
|
|
vec![and(
|
|
binary_expr(col("a"), Operator::Lt, lit(20)),
|
|
binary_expr(col("a"), Operator::Lt, lit(50)),
|
|
)
|
|
.into()], // a < 20 AND a < 50
|
|
vec![0, 1],
|
|
);
|
|
|
|
// test failed to find regions by contradictory filters
|
|
let regions = partition_manager.find_regions_by_filters(
|
|
partition_rule,
|
|
vec![and(
|
|
binary_expr(col("a"), Operator::Lt, lit(20)),
|
|
binary_expr(col("a"), Operator::GtEq, lit(20)),
|
|
)
|
|
.into()]
|
|
.as_slice(),
|
|
); // a < 20 AND a >= 20
|
|
assert!(matches!(
|
|
regions.unwrap_err(),
|
|
partition::error::Error::FindRegions { .. }
|
|
));
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct MockCollector {
|
|
pub write_sum: AtomicU32,
|
|
}
|
|
|
|
impl Collect for MockCollector {
|
|
fn on_write(&self, record: WriteRecord) {
|
|
let _ = self
|
|
.write_sum
|
|
.fetch_add(record.byte_count, Ordering::Relaxed);
|
|
}
|
|
|
|
fn on_read(&self, _record: ReadRecord) {
|
|
todo!()
|
|
}
|
|
}
|
|
|
|
struct MockCalculator;
|
|
|
|
impl WriteCalculator<InsertRequest> for MockCalculator {
|
|
fn calc_byte(&self, _value: &InsertRequest) -> u32 {
|
|
1024 * 10
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn test_meter_insert_request() {
|
|
let collector = Arc::new(MockCollector::default());
|
|
global_registry().set_collector(collector.clone());
|
|
global_registry().register_calculator(Arc::new(MockCalculator));
|
|
|
|
let req = InsertRequest {
|
|
catalog_name: "greptime".to_string(),
|
|
schema_name: "public".to_string(),
|
|
table_name: "numbers".to_string(),
|
|
columns_values: Default::default(),
|
|
region_number: 0,
|
|
};
|
|
meter_insert_request!(req);
|
|
|
|
let re = collector.write_sum.load(Ordering::Relaxed);
|
|
assert_eq!(re, 1024 * 10);
|
|
}
|
|
}
|