rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-13 07:22:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(operator): restore sparse partition tags when rebuilding writes

Browse Source 
Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
WenyXu
2026-10-13 06:37:17 +00:00
committed by GitHub
parent d86741d0ad
commit 9645a03271
4 changed files with 401 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -10135,6 +10135,7 @@ dependencies = [
"meta-client",
"meter-core",
"meter-macros",
"mito-codec",
"moka",
"object-store",
"partition",
src/operator/Cargo.toml
+1
View File
@@ -62,6 +62,7 @@ lazy_static.workspace = true
meta-client.workspace = true
meter-core.workspace = true
meter-macros.workspace = true
mito-codec.workspace = true
moka = { workspace = true, features = ["future"] }
object-store.workspace = true
partition.workspace = true
src/operator/src/req_convert/common/partitioner.rs
+393 -7
View File
@@ -13,6 +13,7 @@
// limitations under the License.
use std::collections::HashMap;
use std::sync::Arc;
use api::v1::region::{
BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest,
@@ -21,12 +22,17 @@ use api::v1::region::{
use api::v1::{ArrowIpc, PartitionExprVersion, Rows};
use common_error::ext::BoxedError;
use common_grpc::flight::{self, FlightDecoder};
use datatypes::arrow::array::{Array, BinaryArray, StringBuilder};
use datatypes::arrow::compute;
use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::arrow::record_batch::RecordBatch;
use mito_codec::row_converter::sparse::{SparseOffsetsCache, SparsePrimaryKeyView};
use partition::PartitionRuleRef;
use partition::manager::PartitionRuleManager;
use partition::partition::RegionMask;
use partition::splitter::RowSplitter;
use snafu::{OptionExt, ResultExt};
use snafu::{OptionExt, ResultExt, ensure};
use store_api::storage::consts;
use store_api::storage::{RegionId, RegionNumber, TableId};
use table::metadata::TableInfo;
@@ -101,9 +107,11 @@ impl TablePartitioner {
///
/// Takes ownership so row payloads can move into their new regions. Requests whose
/// region and partition version still match retain their original payloads.
/// The table metadata comes from the same refresh that created this partitioner.
pub(crate) fn rebuild_write_requests(
&self,
request: RegionRequest,
table_info: &TableInfo,
) -> Result<Vec<RegionRequest>> {
let RegionRequest { header, body } = request;
let bodies = match body.context(error::UnexpectedSnafu {
@@ -111,7 +119,7 @@ impl TablePartitioner {
})? {
Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?,
Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?,
Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk)?,
Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk, table_info)?,
_ => {
return error::UnexpectedSnafu {
violated: "Only writes can be rebuilt",
@@ -165,13 +173,23 @@ impl TablePartitioner {
pub fn partition_bulk_insert_requests(
&self,
batch: RecordBatch,
mut original_ipc: Option<ArrowIpc>,
original_ipc: Option<ArrowIpc>,
skip_wal: bool,
) -> Result<Vec<BulkInsertRequest>> {
let masks = self
.rule
.split_record_batch(&batch)
.context(SplitInsertSnafu)?;
self.build_bulk_insert_requests(batch, masks, original_ipc, skip_wal)
}
fn build_bulk_insert_requests(
&self,
batch: RecordBatch,
masks: HashMap<RegionNumber, RegionMask>,
mut original_ipc: Option<ArrowIpc>,
skip_wal: bool,
) -> Result<Vec<BulkInsertRequest>> {
let mut requests = Vec::with_capacity(masks.len());
for (region, mask) in masks {
if mask.select_none() {
@@ -256,7 +274,11 @@ impl TablePartitioner {
Ok(bodies)
}
fn rebuild_bulk_insert_request(&self, bulk: BulkInsertRequest) -> Result<Vec<Body>> {
fn rebuild_bulk_insert_request(
&self,
bulk: BulkInsertRequest,
table_info: &TableInfo,
) -> Result<Vec<Body>> {
let region_id = RegionId::from_u64(bulk.region_id);
if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) {
return Ok(vec![Body::BulkInsert(bulk)]);
@@ -269,13 +291,92 @@ impl TablePartitioner {
.and_then(|mut decoder| decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload))
.map_err(BoxedError::new)
.context(error::ExternalSnafu)?;
let routing_batch = self.restore_partition_tags(&batch, table_info)?;
let masks = self
.rule
.split_record_batch(routing_batch.as_ref().unwrap_or(&batch))
.context(SplitInsertSnafu)?;
Ok(self
.partition_bulk_insert_requests(batch, Some(ipc), bulk.skip_wal)?
.build_bulk_insert_requests(batch, masks, Some(ipc), bulk.skip_wal)?
.into_iter()
.map(Body::BulkInsert)
.collect())
}
/// Restores omitted physical-table tags for routing without changing the write payload.
fn restore_partition_tags(
&self,
batch: &RecordBatch,
table_info: &TableInfo,
) -> Result<Option<RecordBatch>> {
let meta = &table_info.meta;
if !table_info.is_physical_table() || meta.column_ids.len() != meta.schema.num_columns() {
return Ok(None);
}
let missing = self
.rule
.partition_columns()
.iter()
.filter(|name| batch.column_by_name(name).is_none())
.filter_map(|name| {
let index = meta.schema.column_index_by_name(name)?;
(meta.primary_key_indices.contains(&index)
&& meta.schema.column_schemas()[index].data_type.is_string())
.then_some((name, meta.column_ids[index]))
})
.collect::<Vec<_>>();
if missing.is_empty() {
return Ok(None);
}
let Some(primary_key) = batch.column_by_name(consts::PRIMARY_KEY_COLUMN_NAME) else {
return Ok(None);
};
let primary_key =
primary_key
.as_any()
.downcast_ref::<BinaryArray>()
.context(error::UnexpectedSnafu {
violated: "Sparse primary key column must be Binary",
})?;
let mut builders = missing
.iter()
.map(|_| StringBuilder::new())
.collect::<Vec<_>>();
let mut cache = SparseOffsetsCache::new();
let mut buffer = Vec::new();
for row in 0..batch.num_rows() {
ensure!(
!primary_key.is_null(row),
error::UnexpectedSnafu {
violated: "Sparse primary key must not be null",
}
);
let mut key = SparsePrimaryKeyView::new(primary_key.value(row), &mut cache)
.map_err(BoxedError::new)
.context(error::ExternalSnafu)?;
for ((_, column_id), builder) in missing.iter().zip(&mut builders) {
let label = key
.label(*column_id, &mut buffer)
.map_err(BoxedError::new)
.context(error::ExternalSnafu)?;
builder.append_option(label);
}
}
let schema = batch.schema();
let mut fields = schema.fields().to_vec();
let mut columns = batch.columns().to_vec();
for ((name, _), mut builder) in missing.into_iter().zip(builders) {
fields.push(Arc::new(Field::new(name, DataType::Utf8, true)));
columns.push(Arc::new(builder.finish()));
}
RecordBatch::try_new(
Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())),
columns,
)
.map(Some)
.context(error::ComputeArrowSnafu)
}
/// Whether the request's region still exists with the same partition expression version.
///
/// A leader change leaves the partition unchanged, so the original payload can be
@@ -313,17 +414,302 @@ mod tests {
use common_meta::key::table_route::TableRouteValue;
use common_meta::kv_backend::memory::MemoryKvBackend;
use common_meta::kv_backend::test_util::MockKvBackendBuilder;
use common_meta::rpc::router::{Region, RegionRoute};
use common_meta::rpc::store::RangeResponse;
use datatypes::arrow::array::Int32Array;
use datatypes::arrow::array::{Int32Array, StringArray};
use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::data_type::ConcreteDataType;
use datatypes::schema;
use datatypes::value::{Value as DataValue, ValueRef};
use mito_codec::row_converter::sparse;
use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec};
use moka::future::Cache;
use partition::expr;
use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
use partition::multi_dim::MultiDimPartitionRule;
use store_api::metric_engine_consts;
use super::*;
use crate::error::Error;
use crate::test_util;
fn build_sparse_partitioner() -> TablePartitioner {
TablePartitioner {
table_id: 1,
rule: Arc::new(
MultiDimPartitionRule::try_new(
vec!["zone".to_string()],
vec![1, 2],
vec![
expr::col("zone").lt(DataValue::String("m".into())),
expr::col("zone").gt_eq(DataValue::String("m".into())),
],
true,
)
.unwrap(),
),
versions: [
(1, Some(PartitionExprVersion { value: 2 })),
(2, Some(PartitionExprVersion { value: 2 })),
]
.into_iter()
.collect(),
}
}
fn build_sparse_table_info() -> TableInfo {
let mut table_info = test_util::new_test_table_info(1, "physical_tags", [1, 2].into_iter());
table_info.meta.schema = Arc::new(schema::Schema::new(vec![
schema::ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
schema::ColumnSchema::new("zone", ConcreteDataType::string_datatype(), true),
]));
table_info.meta.primary_key_indices = vec![0, 1];
table_info.meta.partition_key_indices = vec![1];
table_info.meta.column_ids = vec![10, 20];
table_info.meta.options.extra_options.insert(
metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(),
String::new(),
);
table_info
}
fn build_sparse_batch() -> RecordBatch {
let codec = SparsePrimaryKeyCodec::schemaless();
let keys = [Some("east"), Some("west"), Some(""), None, None]
.into_iter()
.enumerate()
.map(|(index, zone)| {
let mut key = Vec::new();
let values = [
(
sparse::RESERVED_COLUMN_ID_TABLE_ID,
ValueRef::UInt32(42 + index as u32),
),
(
sparse::RESERVED_COLUMN_ID_TSID,
ValueRef::UInt64(index as u64),
),
(10, ValueRef::String("node")),
];
codec.encode_value_refs(&values, &mut key).unwrap();
if let Some(zone) = zone {
codec
.encode_raw_tag_value([(20, zone.as_bytes())].into_iter(), &mut key)
.unwrap();
}
if index == 4 {
key.extend_from_slice(&20_u32.to_be_bytes());
key.push(0);
}
key
})
.collect::<Vec<_>>();
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("host", DataType::Utf8, false),
Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false),
])),
vec![
Arc::new(StringArray::from(vec!["node"; keys.len()])),
Arc::new(BinaryArray::from_iter_values(keys.iter())),
],
)
.unwrap()
}
#[tokio::test]
async fn test_partition_tags_use_physical_table_metadata() {
let backend = Arc::new(MemoryKvBackend::new());
let manager = test_util::create_partition_rule_manager(backend.clone()).await;
let table_info = build_sparse_table_info();
let routes = [
expr::col("zone").lt(DataValue::String("m".into())),
expr::col("zone").gt_eq(DataValue::String("m".into())),
]
.into_iter()
.enumerate()
.map(|(index, expr)| RegionRoute {
region: Region {
id: RegionId::new(1, index as u32 + 1),
partition_expr: expr.as_json_str().unwrap(),
..Default::default()
},
..Default::default()
})
.collect();
TableMetadataManager::new(backend)
.create_table_metadata(
table_info.clone(),
TableRouteValue::physical(routes),
Default::default(),
)
.await
.unwrap();
let batch = build_sparse_batch();
let partitioner = TablePartitioner::try_new(&manager, &table_info)
.await
.unwrap();
let restored = partitioner
.restore_partition_tags(&batch, &table_info)
.unwrap()
.unwrap();
let zones = restored
.column_by_name("zone")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
assert_eq!(
zones.iter().collect::<Vec<_>>(),
vec![Some("east"), Some("west"), Some(""), None, None]
);
for (physical, tag, valid_ids) in [
(false, true, true),
(true, false, true),
(true, true, false),
] {
let mut info = table_info.clone();
if !physical {
info.meta
.options
.extra_options
.remove(metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY);
}
if !tag {
info.meta.primary_key_indices = vec![0];
}
if !valid_ids {
info.meta.column_ids.clear();
}
let partitioner = TablePartitioner::try_new(&manager, &info).await.unwrap();
assert!(
partitioner
.restore_partition_tags(&batch, &info)
.unwrap()
.is_none()
);
assert!(
partitioner
.partition_bulk_insert_requests(batch.clone(), None, false)
.is_err()
);
}
}
#[test]
fn test_rebuild_sparse_bulk_with_changed_partition_column() {
let partitioner = build_sparse_partitioner();
let table_info = build_sparse_table_info();
let batch = build_sparse_batch();
let restored = partitioner
.restore_partition_tags(&batch, &table_info)
.unwrap()
.unwrap();
let zones = restored
.column_by_name("zone")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
assert_eq!(
zones.iter().collect::<Vec<_>>(),
vec![Some("east"), Some("west"), Some(""), None, None]
);
let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap();
let rebuilt = partitioner
.rebuild_write_requests(
build_request(Body::BulkInsert(BulkInsertRequest {
region_id: RegionId::new(1, 1).as_u64(),
partition_expr_version: Some(PartitionExprVersion { value: 1 }),
skip_wal: true,
body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
schema,
data_header,
payload,
})),
..Default::default()
})),
&table_info,
)
.unwrap();
assert_eq!(rebuilt.len(), 2);
for request in rebuilt {
let Some(Body::BulkInsert(bulk)) = request.body else {
unreachable!()
};
assert!(bulk.skip_wal);
assert_eq!(
bulk.partition_expr_version,
Some(PartitionExprVersion { value: 2 })
);
let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else {
unreachable!()
};
let actual = FlightDecoder::try_from_schema_bytes(&ipc.schema)
.unwrap()
.try_decode_record_batch(&ipc.data_header, &ipc.payload)
.unwrap();
let mask = if RegionId::from_u64(bulk.region_id).region_number() == 1 {
vec![true, false, true, true, true]
} else {
vec![false, true, false, false, false]
};
let expected = compute::filter_record_batch(&batch, &mask.into()).unwrap();
assert_eq!(actual, expected);
}
}
#[test]
fn test_restore_partition_tags_checks_only_missing_tags() {
let partitioner = build_sparse_partitioner();
let table_info = build_sparse_table_info();
let batch = build_sparse_batch();
let invalid = RecordBatch::try_new(
batch.schema(),
vec![
batch.column(0).clone(),
Arc::new(BinaryArray::from_iter_values([b"invalid"; 5])),
],
)
.unwrap();
assert!(
partitioner
.restore_partition_tags(&invalid, &table_info)
.is_err()
);
let mut fields = invalid.schema().fields().to_vec();
fields.push(Arc::new(Field::new("zone", DataType::Utf8, true)));
let mut columns = invalid.columns().to_vec();
columns.push(Arc::new(StringArray::from(vec![Some("east"); 5])));
let complete = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap();
assert!(
partitioner
.restore_partition_tags(&complete, &table_info)
.unwrap()
.is_none()
);
let (schema, data_header, payload) = flight::record_batch_to_ipc(invalid).unwrap();
let request = build_request(Body::BulkInsert(BulkInsertRequest {
region_id: RegionId::new(1, 1).as_u64(),
partition_expr_version: Some(PartitionExprVersion { value: 2 }),
body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
schema,
data_header,
payload,
})),
..Default::default()
}));
let expected = request.clone();
assert_eq!(
partitioner
.rebuild_write_requests(request, &table_info)
.unwrap(),
vec![expected]
);
}
#[tokio::test]
async fn test_partition_requests_preserve_routing_errors() {
let table_info = test_util::new_test_table_info(1, "table_1", [1].into_iter());
@@ -693,6 +1079,6 @@ mod tests {
let partitioner = TablePartitioner::try_new(manager, &table_info)
.await
.context(error::InvalidPartitionSnafu)?;
partitioner.rebuild_write_requests(request)
partitioner.rebuild_write_requests(request, &table_info)
}
}
src/operator/src/write_dispatcher.rs
+6 -5
View File
@@ -390,7 +390,7 @@ impl WriteRetry {
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
partitioners.insert(table_id, partitioner);
partitioners.insert(table_id, (table_info, partitioner));
}
let mut rebuilt_requests = Vec::new();
for write in writes {
@@ -401,11 +401,12 @@ impl WriteRetry {
.build()
})?;
let table_id = write_region_id(&request)?.table_id();
let partitioner = partitioners.get(&table_id).context(UnexpectedSnafu {
err_msg: "Refreshed write has no partition snapshot",
})?;
let (table_info, partitioner) =
partitioners.get(&table_id).context(UnexpectedSnafu {
err_msg: "Refreshed write has no partition snapshot",
})?;
let requests = partitioner
.rebuild_write_requests(request)
.rebuild_write_requests(request, table_info)
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
rebuilt_requests.extend(requests);