diff --git a/src/operator/src/tests/partition_manager.rs b/src/operator/src/tests/partition_manager.rs index b884a1d951..ecf87aedac 100644 --- a/src/operator/src/tests/partition_manager.rs +++ b/src/operator/src/tests/partition_manager.rs @@ -15,6 +15,9 @@ use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; +use api::v1::helper::tag_column_schema; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Row, Rows}; use common_base::hash::partition_expr_version; use common_meta::cache::{TableRouteCacheRef, new_table_route_cache}; use common_meta::key::TableMetadataManager; @@ -37,7 +40,11 @@ pub fn new_test_table_info( table_name: &str, _region_numbers: impl Iterator, ) -> TableInfo { - let column_schemas = vec![ + new_test_table_info_with_columns(table_id, table_name, test_column_schemas(true), vec![0]) +} + +fn test_column_schemas(include_b: bool) -> Vec { + let mut column_schemas = vec![ ColumnSchema::new("a", ConcreteDataType::int32_datatype(), true), ColumnSchema::new( "ts", @@ -45,8 +52,24 @@ pub fn new_test_table_info( false, ) .with_time_index(true), - ColumnSchema::new("b", ConcreteDataType::int32_datatype(), true), ]; + if include_b { + column_schemas.push(ColumnSchema::new( + "b", + ConcreteDataType::int32_datatype(), + true, + )); + } + column_schemas +} + +fn new_test_table_info_with_columns( + table_id: u32, + table_name: &str, + column_schemas: Vec, + partition_key_indices: Vec, +) -> TableInfo { + let next_column_id = column_schemas.len() as u32; let schema = SchemaBuilder::try_from(column_schemas) .unwrap() .version(123) @@ -55,10 +78,10 @@ pub fn new_test_table_info( let meta = TableMetaBuilder::empty() .schema(Arc::new(schema)) - .primary_key_indices(vec![0]) + .primary_key_indices(partition_key_indices.clone()) .engine("engine") - .next_column_id(3) - .partition_key_indices(vec![0]) + .next_column_id(next_column_id) + .partition_key_indices(partition_key_indices) .build() .unwrap(); TableInfoBuilder::default() @@ -70,6 +93,14 @@ pub fn new_test_table_info( .unwrap() } +fn new_physical_test_table_info(table_id: u32, table_name: &str) -> TableInfo { + new_test_table_info_with_columns(table_id, table_name, test_column_schemas(true), vec![0, 2]) +} + +fn new_logical_test_table_info(table_id: u32, table_name: &str) -> TableInfo { + new_test_table_info_with_columns(table_id, table_name, test_column_schemas(false), vec![0]) +} + fn new_test_region_wal_options(regions: Vec) -> RegionWalOptions { // TODO(niebayes): construct region wal options for test. let _ = regions; @@ -239,3 +270,97 @@ async fn test_partition_expr_version_cache() { assert_ne!(None, *version_by_region.get(&2).unwrap()); assert_ne!(None, *version_by_region.get(&3).unwrap()); } + +#[tokio::test] +async fn test_logical_split_rows_rejects_physical_only_partition_column() { + let kv_backend = Arc::new(common_meta::kv_backend::memory::MemoryKvBackend::new()); + let table_metadata_manager = TableMetadataManager::new(kv_backend.clone()); + let table_route_cache = test_new_table_route_cache(kv_backend.clone()); + let partition_info_cache = test_new_partition_info_cache(table_route_cache.clone()); + let partition_manager = + PartitionRuleManager::new(kv_backend, table_route_cache, partition_info_cache); + let physical_table_info = new_physical_test_table_info(1024, "physical"); + let logical_table_info = new_logical_test_table_info(1025, "logical"); + assert_eq!( + vec!["a", "b"], + physical_table_info + .meta + .partition_column_names() + .map(|name| name.as_str()) + .collect::>() + ); + assert_eq!( + vec!["a"], + logical_table_info + .meta + .partition_column_names() + .map(|name| name.as_str()) + .collect::>() + ); + + table_metadata_manager + .create_table_metadata( + physical_table_info, + TableRouteValue::physical(vec![RegionRoute { + region: Region { + id: 1.into(), + name: "r1".to_string(), + attrs: BTreeMap::new(), + partition_expr: PartitionExpr::new( + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::GtEq, + Operand::Value(datatypes::value::Value::Int32(10)), + )), + RestrictedOp::And, + Operand::Expr(PartitionExpr::new( + Operand::Column("b".to_string()), + RestrictedOp::Lt, + Operand::Value(datatypes::value::Value::Int32(50)), + )), + ) + .as_json_str() + .unwrap(), + }, + leader_peer: Some(Peer::new(1, "")), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + new_test_region_wal_options(vec![1]), + ) + .await + .unwrap(); + table_metadata_manager + .create_table_metadata( + logical_table_info.clone(), + TableRouteValue::logical(1024), + new_test_region_wal_options(vec![]), + ) + .await + .unwrap(); + + let (partition_rule, _) = partition_manager + .find_table_partition_rule(&logical_table_info) + .await + .unwrap(); + assert_eq!(&["a"], partition_rule.partition_columns()); + + let error = partition_manager + .split_rows( + &logical_table_info, + Rows { + schema: vec![tag_column_schema("a", ColumnDataType::Int32)], + rows: vec![Row { + values: vec![ValueData::I32Value(100).into()], + }], + }, + ) + .await + .unwrap_err(); + assert!(matches!( + error, + partition::error::Error::UndefinedColumn { ref column, .. } if column == "b" + )); +} diff --git a/src/partition/src/multi_dim.rs b/src/partition/src/multi_dim.rs index a68479888b..df7b5291e9 100644 --- a/src/partition/src/multi_dim.rs +++ b/src/partition/src/multi_dim.rs @@ -121,12 +121,18 @@ impl MultiDimPartitionRule { fn evaluate_expr(&self, expr: &PartitionExpr, values: &[Value]) -> Result { match (expr.lhs.as_ref(), expr.rhs.as_ref()) { (Operand::Column(name), Operand::Value(r)) => { - let index = self.name_to_index.get(name).unwrap(); + let index = self + .name_to_index + .get(name) + .context(UndefinedColumnSnafu { column: name })?; let l = &values[*index]; Self::perform_op(l, &expr.op, r) } (Operand::Value(l), Operand::Column(name)) => { - let index = self.name_to_index.get(name).unwrap(); + let index = self + .name_to_index + .get(name) + .context(UndefinedColumnSnafu { column: name })?; let r = &values[*index]; Self::perform_op(l, &expr.op, r) } @@ -416,6 +422,46 @@ mod tests { assert_matches!(rule.find_region(&["zzzz".into()]), Ok(3)); } + #[test] + fn test_find_region_rejects_undeclared_column_on_lhs() { + let rule = MultiDimPartitionRule::try_new( + vec!["host".to_string()], + vec![1], + vec![PartitionExpr::new( + Operand::Column("rack".to_string()), + RestrictedOp::Lt, + Operand::Value(Value::String("n".into())), + )], + false, + ) + .unwrap(); + + assert_matches!( + rule.find_region(&[Value::String("z".into())]), + Err(Error::UndefinedColumn { column, .. }) if column == "rack" + ); + } + + #[test] + fn test_find_region_rejects_undeclared_column_on_rhs() { + let rule = MultiDimPartitionRule::try_new( + vec!["host".to_string()], + vec![1], + vec![PartitionExpr::new( + Operand::Value(Value::String("n".into())), + RestrictedOp::Gt, + Operand::Column("rack".to_string()), + )], + false, + ) + .unwrap(); + + assert_matches!( + rule.find_region(&[Value::String("z".into())]), + Err(Error::UndefinedColumn { column, .. }) if column == "rack" + ); + } + #[test] fn invalid_expr_case_1() { // PARTITION ON COLUMNS (b) (