meta_srv/procedure/repartition/

plan.rs

// 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::cmp::Ordering;

use common_meta::rpc::router::RegionRoute;
use partition::expr::PartitionExpr;
use serde::{Deserialize, Deserializer, Serialize};
use snafu::ResultExt;
use store_api::storage::{RegionId, RegionNumber, TableId};

use crate::error::{self, Result};
use crate::procedure::repartition::group::GroupId;

/// Metadata describing a source region involved in the plan.
///
/// Source regions may represent either an existing partitioned region or the
/// default region of an unpartitioned table.
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub enum SourceRegionDescriptor {
    /// A regular partitioned source region.
    Partitioned {
        /// The region id of the source region.
        region_id: RegionId,
        /// The partition expression of the source region.
        partition_expr: PartitionExpr,
    },
    /// The default source region of an unpartitioned table.
    Default {
        /// The region id of the default source region.
        region_id: RegionId,
    },
}

impl<'de> Deserialize<'de> for SourceRegionDescriptor {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(deny_unknown_fields)]
        struct PartitionedSourceRegionDescriptor {
            region_id: RegionId,
            partition_expr: PartitionExpr,
        }

        #[derive(Deserialize)]
        #[serde(untagged)]
        enum SourceRegionDescriptorRepr {
            Tagged(SourceRegionDescriptorTaggedRepr),
            Legacy(PartitionedSourceRegionDescriptor),
        }

        #[derive(Deserialize)]
        enum SourceRegionDescriptorTaggedRepr {
            Partitioned {
                region_id: RegionId,
                partition_expr: PartitionExpr,
            },
            Default {
                region_id: RegionId,
            },
        }

        match SourceRegionDescriptorRepr::deserialize(deserializer)? {
            SourceRegionDescriptorRepr::Tagged(SourceRegionDescriptorTaggedRepr::Partitioned {
                region_id,
                partition_expr,
            }) => Ok(Self::Partitioned {
                region_id,
                partition_expr,
            }),
            SourceRegionDescriptorRepr::Tagged(SourceRegionDescriptorTaggedRepr::Default {
                region_id,
            }) => Ok(Self::Default { region_id }),
            SourceRegionDescriptorRepr::Legacy(descriptor) => Ok(Self::Partitioned {
                region_id: descriptor.region_id,
                partition_expr: descriptor.partition_expr,
            }),
        }
    }
}

impl SourceRegionDescriptor {
    /// Creates a partitioned source region descriptor.
    pub fn partitioned(region_id: RegionId, partition_expr: PartitionExpr) -> Self {
        Self::Partitioned {
            region_id,
            partition_expr,
        }
    }

    /// Returns the region id of this source descriptor.
    pub fn region_id(&self) -> RegionId {
        match self {
            Self::Partitioned { region_id, .. } => *region_id,
            Self::Default { region_id } => *region_id,
        }
    }

    /// Returns the partition expression if this source is partitioned.
    pub fn partition_expr(&self) -> Option<&PartitionExpr> {
        match self {
            Self::Partitioned { partition_expr, .. } => Some(partition_expr),
            Self::Default { .. } => None,
        }
    }

    /// Returns true if this source descriptor matches the route partition expression.
    pub fn matches_route_expr(&self, route_expr: &str) -> Result<bool> {
        match self {
            Self::Partitioned { partition_expr, .. } => {
                let expected = partition_expr
                    .as_json_str()
                    .context(error::SerializePartitionExprSnafu)?;
                Ok(route_expr == expected)
            }
            Self::Default { .. } => Ok(route_expr.is_empty()),
        }
    }

    /// Returns the route partition expression to restore during rollback.
    pub fn route_expr_for_rollback(&self) -> Result<String> {
        match self {
            Self::Partitioned { partition_expr, .. } => partition_expr
                .as_json_str()
                .context(error::SerializePartitionExprSnafu),
            Self::Default { .. } => Ok(String::new()),
        }
    }
}

/// Metadata describing a target region involved in the plan.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct TargetRegionDescriptor {
    /// The region id of the target region.
    pub region_id: RegionId,
    /// The partition expression of the target region.
    pub partition_expr: PartitionExpr,
}

/// A plan entry for the region allocation phase, describing source regions
/// and target partition expressions before allocation.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct AllocationPlanEntry {
    /// The group id for this plan entry.
    pub group_id: GroupId,
    /// Source region descriptors involved in the plan.
    pub source_regions: Vec<SourceRegionDescriptor>,
    /// The target partition expressions for the new or changed regions.
    pub target_partition_exprs: Vec<PartitionExpr>,
    /// For each `source_regions[k]`, the corresponding vector contains global
    /// `target_partition_exprs` that overlap with it.
    pub transition_map: Vec<Vec<usize>>,
}

/// A plan entry for the dispatch phase after region allocation,
/// with concrete source and target region descriptors.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RepartitionPlanEntry {
    /// The group id for this plan entry.
    pub group_id: GroupId,
    /// The source region descriptors involved in the plan.
    pub source_regions: Vec<SourceRegionDescriptor>,
    /// The target region descriptors involved in the plan.
    pub target_regions: Vec<TargetRegionDescriptor>,
    /// The region ids of the allocated regions.
    pub allocated_region_ids: Vec<RegionId>,
    /// The region ids of the regions that are pending deallocation.
    pub pending_deallocate_region_ids: Vec<RegionId>,
    /// For each `source_regions[k]`, the corresponding vector contains global
    /// `target_regions` that overlap with it.
    pub transition_map: Vec<Vec<usize>>,
    /// Pre-staging target routes persisted for parent rollback and recovery.
    #[serde(default)]
    pub original_target_routes: Vec<RegionRoute>,
}

impl RepartitionPlanEntry {
    /// Returns the target regions that are newly allocated.
    pub(crate) fn allocate_regions(&self) -> Vec<&TargetRegionDescriptor> {
        self.target_regions
            .iter()
            .filter(|r| self.allocated_region_ids.contains(&r.region_id))
            .collect()
    }
}

/// Converts an allocation plan to a repartition plan.
///
/// Converts an [`AllocationPlanEntry`] (which contains abstract region allocation intents)
/// into a [`RepartitionPlanEntry`] with concrete source and target region descriptors,
/// plus records information on which regions are newly allocated and/or pending deallocation.
///
/// # Returns
///
/// A [`Result`] containing the [`RepartitionPlanEntry`] describing exactly the source regions,
/// the target regions (including any that need to be newly allocated), and transition mappings.
///
/// # Notes
/// - If new regions are needed, their region ids are constructed using `table_id` and incrementing
///   from `next_region_number`.
/// - For each target, associates the correct region descriptor; for new regions, the region id
///   is assigned sequentially.
pub fn convert_allocation_plan_to_repartition_plan(
    table_id: TableId,
    next_region_number: &mut RegionNumber,
    AllocationPlanEntry {
        group_id,
        source_regions,
        target_partition_exprs,
        transition_map,
        ..
    }: &AllocationPlanEntry,
) -> RepartitionPlanEntry {
    match source_regions.len().cmp(&target_partition_exprs.len()) {
        Ordering::Less => {
            // requires to allocate regions
            let pending_allocate_target_partition_exprs = target_partition_exprs
                .iter()
                .skip(source_regions.len())
                .map(|target_partition_expr| {
                    let desc = TargetRegionDescriptor {
                        region_id: RegionId::new(table_id, *next_region_number),
                        partition_expr: target_partition_expr.clone(),
                    };
                    *next_region_number += 1;
                    desc
                })
                .collect::<Vec<_>>();

            let allocated_region_ids = pending_allocate_target_partition_exprs
                .iter()
                .map(|rd| rd.region_id)
                .collect::<Vec<_>>();

            let target_regions = source_regions
                .iter()
                .zip(target_partition_exprs.iter())
                .map(
                    |(source_region, target_partition_expr)| TargetRegionDescriptor {
                        region_id: source_region.region_id(),
                        partition_expr: target_partition_expr.clone(),
                    },
                )
                .chain(pending_allocate_target_partition_exprs)
                .collect::<Vec<_>>();

            RepartitionPlanEntry {
                group_id: *group_id,
                source_regions: source_regions.clone(),
                target_regions,
                allocated_region_ids,
                pending_deallocate_region_ids: vec![],
                transition_map: transition_map.clone(),
                original_target_routes: vec![],
            }
        }
        Ordering::Equal => {
            let target_regions = source_regions
                .iter()
                .zip(target_partition_exprs.iter())
                .map(
                    |(source_region, target_partition_expr)| TargetRegionDescriptor {
                        region_id: source_region.region_id(),
                        partition_expr: target_partition_expr.clone(),
                    },
                )
                .collect::<Vec<_>>();

            RepartitionPlanEntry {
                group_id: *group_id,
                source_regions: source_regions.clone(),
                target_regions,
                allocated_region_ids: vec![],
                pending_deallocate_region_ids: vec![],
                transition_map: transition_map.clone(),
                original_target_routes: vec![],
            }
        }
        Ordering::Greater => {
            // requires to deallocate regions
            let target_regions = source_regions
                .iter()
                .take(target_partition_exprs.len())
                .zip(target_partition_exprs.iter())
                .map(
                    |(source_region, target_partition_expr)| TargetRegionDescriptor {
                        region_id: source_region.region_id(),
                        partition_expr: target_partition_expr.clone(),
                    },
                )
                .collect::<Vec<_>>();

            let pending_deallocate_region_ids = source_regions
                .iter()
                .skip(target_partition_exprs.len())
                .map(|source_region| source_region.region_id())
                .collect::<Vec<_>>();

            RepartitionPlanEntry {
                group_id: *group_id,
                source_regions: source_regions.clone(),
                target_regions,
                allocated_region_ids: vec![],
                pending_deallocate_region_ids,
                transition_map: transition_map.clone(),
                original_target_routes: vec![],
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use store_api::storage::RegionId;
    use uuid::Uuid;

    use super::*;
    use crate::procedure::repartition::test_util::range_expr;

    fn create_region_descriptor(
        table_id: TableId,
        region_number: u32,
        col: &str,
        start: i64,
        end: i64,
    ) -> SourceRegionDescriptor {
        SourceRegionDescriptor::partitioned(
            RegionId::new(table_id, region_number),
            range_expr(col, start, end),
        )
    }

    #[test]
    fn test_source_region_descriptor_deserializes_legacy_partitioned_shape() {
        let table_id = 1024;
        let region_id = RegionId::new(table_id, 1);
        let partition_expr = range_expr("x", 0, 100);
        let legacy_json = serde_json::json!({
            "region_id": region_id,
            "partition_expr": partition_expr,
        });

        let descriptor: SourceRegionDescriptor = serde_json::from_value(legacy_json).unwrap();

        assert_eq!(
            descriptor,
            SourceRegionDescriptor::partitioned(region_id, partition_expr)
        );
    }

    #[test]
    fn test_source_region_descriptor_rejects_legacy_default_shape() {
        let region_id = RegionId::new(1024, 1);
        let default_json = serde_json::json!({
            "region_id": region_id,
        });

        let err = serde_json::from_value::<SourceRegionDescriptor>(default_json).unwrap_err();

        assert!(err.to_string().contains("data did not match any variant"));
    }

    #[test]
    fn test_source_region_descriptor_roundtrip_tagged_partitioned_shape() {
        let region_id = RegionId::new(1024, 1);
        let partition_expr = range_expr("x", 0, 100);
        let descriptor = SourceRegionDescriptor::partitioned(region_id, partition_expr.clone());

        let value = serde_json::to_value(&descriptor).unwrap();
        let decoded = serde_json::from_value::<SourceRegionDescriptor>(value.clone()).unwrap();

        assert_eq!(
            value,
            serde_json::json!({
                "Partitioned": {
                    "region_id": region_id,
                    "partition_expr": partition_expr,
                }
            })
        );
        assert_eq!(decoded, descriptor);
    }

    #[test]
    fn test_source_region_descriptor_roundtrip_tagged_default_shape() {
        let region_id = RegionId::new(1024, 1);
        let descriptor = SourceRegionDescriptor::Default { region_id };

        let value = serde_json::to_value(&descriptor).unwrap();
        let decoded = serde_json::from_value::<SourceRegionDescriptor>(value.clone()).unwrap();

        assert_eq!(
            value,
            serde_json::json!({
                "Default": {
                    "region_id": region_id,
                }
            })
        );
        assert_eq!(decoded, descriptor);
    }

    #[test]
    fn test_source_region_descriptor_rejects_invalid_partition_expr_shape() {
        let region_id = RegionId::new(1024, 1);
        let invalid_json = serde_json::json!({
            "region_id": region_id,
            "partition_expr": 42,
        });

        let err = serde_json::from_value::<SourceRegionDescriptor>(invalid_json).unwrap_err();

        assert!(err.to_string().contains("data did not match any variant"));
    }

    #[test]
    fn test_repartition_plan_entry_deserializes_legacy_source_regions() {
        let group_id = Uuid::new_v4();
        let table_id = 1024;
        let source_region_id = RegionId::new(table_id, 1);
        let target_region_id = RegionId::new(table_id, 2);
        let source_partition_expr = range_expr("x", 0, 100);
        let target_partition_expr = range_expr("x", 0, 50);
        let legacy_json = serde_json::json!({
            "group_id": group_id,
            "source_regions": [{
                "region_id": source_region_id,
                "partition_expr": source_partition_expr,
            }],
            "target_regions": [{
                "region_id": target_region_id,
                "partition_expr": target_partition_expr,
            }],
            "allocated_region_ids": [target_region_id],
            "pending_deallocate_region_ids": [],
            "transition_map": [[0]],
        });

        let plan: RepartitionPlanEntry = serde_json::from_value(legacy_json).unwrap();

        assert_eq!(plan.group_id, group_id);
        assert_eq!(
            plan.source_regions,
            vec![SourceRegionDescriptor::partitioned(
                source_region_id,
                source_partition_expr
            )]
        );
        assert_eq!(
            plan.target_regions,
            vec![TargetRegionDescriptor {
                region_id: target_region_id,
                partition_expr: target_partition_expr,
            }]
        );
        assert_eq!(plan.allocated_region_ids, vec![target_region_id]);
        assert!(plan.pending_deallocate_region_ids.is_empty());
        assert_eq!(plan.transition_map, vec![vec![0]]);
        assert!(plan.original_target_routes.is_empty());
    }

    #[test]
    fn test_convert_plan_equal_regions() {
        let group_id = Uuid::new_v4();
        let table_id = 1024;
        let mut next_region_number = 10;
        let source_regions = vec![
            create_region_descriptor(table_id, 1, "x", 0, 100),
            create_region_descriptor(table_id, 2, "x", 100, 200),
        ];
        let target_partition_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 200)];
        let allocation_plan = AllocationPlanEntry {
            group_id,
            source_regions: source_regions.clone(),
            target_partition_exprs: target_partition_exprs.clone(),
            transition_map: Vec::new(),
        };
        let result = convert_allocation_plan_to_repartition_plan(
            table_id,
            &mut next_region_number,
            &allocation_plan,
        );
        assert_eq!(result.group_id, group_id);
        assert_eq!(result.source_regions, source_regions);
        assert_eq!(result.target_regions.len(), 2);
        assert!(result.allocated_region_ids.is_empty());
        assert!(result.pending_deallocate_region_ids.is_empty());
        // next_region_number should not change for equal regions
        assert_eq!(next_region_number, 10);
        // Verify target regions
        assert_eq!(
            result.target_regions[0].region_id,
            RegionId::new(table_id, 1)
        );
        assert_eq!(
            result.target_regions[0].partition_expr,
            target_partition_exprs[0]
        );
        assert_eq!(
            result.target_regions[1].region_id,
            RegionId::new(table_id, 2)
        );
        assert_eq!(
            result.target_regions[1].partition_expr,
            target_partition_exprs[1]
        );
    }

    #[test]
    fn test_convert_plan_allocate_regions() {
        let group_id = Uuid::new_v4();
        let table_id = 1024;
        let mut next_region_number = 10;

        // 3 source regions -> 5 target partition expressions
        let source_regions = vec![
            create_region_descriptor(table_id, 1, "x", 0, 100),
            create_region_descriptor(table_id, 2, "x", 100, 200),
            create_region_descriptor(table_id, 3, "x", 200, 300),
        ];
        let target_partition_exprs = vec![
            range_expr("x", 0, 50),
            range_expr("x", 50, 100),
            range_expr("x", 100, 150),
            range_expr("x", 150, 200),
            range_expr("x", 200, 300),
        ];
        let allocation_plan = AllocationPlanEntry {
            group_id,
            source_regions: source_regions.clone(),
            target_partition_exprs: target_partition_exprs.clone(),
            transition_map: vec![],
        };
        let result = convert_allocation_plan_to_repartition_plan(
            table_id,
            &mut next_region_number,
            &allocation_plan,
        );
        assert_eq!(result.group_id, group_id);
        assert_eq!(result.source_regions, source_regions);
        assert_eq!(result.target_regions.len(), 5);
        assert_eq!(result.allocated_region_ids.len(), 2);
        assert!(result.pending_deallocate_region_ids.is_empty());
        assert_eq!(next_region_number, 12);

        // Verify first 3 target regions use source region ids with target partition exprs
        assert_eq!(
            result.target_regions[0].region_id,
            RegionId::new(table_id, 1)
        );
        assert_eq!(
            result.target_regions[0].partition_expr,
            target_partition_exprs[0]
        );
        assert_eq!(
            result.target_regions[1].region_id,
            RegionId::new(table_id, 2)
        );
        assert_eq!(
            result.target_regions[1].partition_expr,
            target_partition_exprs[1]
        );
        assert_eq!(
            result.target_regions[2].region_id,
            RegionId::new(table_id, 3)
        );
        assert_eq!(
            result.target_regions[2].partition_expr,
            target_partition_exprs[2]
        );

        // Verify last 2 target regions are newly allocated
        assert_eq!(
            result.target_regions[3].region_id,
            RegionId::new(table_id, 10)
        );
        assert_eq!(
            result.target_regions[3].partition_expr,
            target_partition_exprs[3]
        );
        assert_eq!(
            result.target_regions[4].region_id,
            RegionId::new(table_id, 11)
        );
        assert_eq!(
            result.target_regions[4].partition_expr,
            target_partition_exprs[4]
        );

        // Verify allocated region ids
        assert_eq!(result.allocated_region_ids[0], RegionId::new(table_id, 10));
        assert_eq!(result.allocated_region_ids[1], RegionId::new(table_id, 11));
    }

    #[test]
    fn test_convert_plan_deallocate_regions() {
        let group_id = Uuid::new_v4();
        let table_id = 1024;

        // 5 source regions -> 3 target partition expressions
        let source_regions = vec![
            create_region_descriptor(table_id, 1, "x", 0, 50),
            create_region_descriptor(table_id, 2, "x", 50, 100),
            create_region_descriptor(table_id, 3, "x", 100, 150),
            create_region_descriptor(table_id, 4, "x", 150, 200),
            create_region_descriptor(table_id, 5, "x", 200, 300),
        ];
        let target_partition_exprs = vec![
            range_expr("x", 0, 100),
            range_expr("x", 100, 200),
            range_expr("x", 200, 300),
        ];
        let allocation_plan = AllocationPlanEntry {
            group_id,
            source_regions: source_regions.clone(),
            target_partition_exprs: target_partition_exprs.clone(),
            transition_map: vec![],
        };
        let mut next_region_number = 10;
        let result = convert_allocation_plan_to_repartition_plan(
            table_id,
            &mut next_region_number,
            &allocation_plan,
        );
        assert_eq!(next_region_number, 10);
        assert_eq!(result.group_id, group_id);
        assert_eq!(result.source_regions, source_regions);
        assert_eq!(result.target_regions.len(), 3);
        assert!(result.allocated_region_ids.is_empty());
        assert_eq!(result.pending_deallocate_region_ids.len(), 2);

        // Verify first 3 source regions are kept as target regions with new partition exprs
        assert_eq!(
            result.target_regions[0].region_id,
            RegionId::new(table_id, 1)
        );
        assert_eq!(
            result.target_regions[0].partition_expr,
            target_partition_exprs[0]
        );
        assert_eq!(
            result.target_regions[1].region_id,
            RegionId::new(table_id, 2)
        );
        assert_eq!(
            result.target_regions[1].partition_expr,
            target_partition_exprs[1]
        );
        assert_eq!(
            result.target_regions[2].region_id,
            RegionId::new(table_id, 3)
        );
        assert_eq!(
            result.target_regions[2].partition_expr,
            target_partition_exprs[2]
        );

        // Verify last 2 source regions are pending deallocation
        assert_eq!(
            result.pending_deallocate_region_ids[0],
            RegionId::new(table_id, 4)
        );
        assert_eq!(
            result.pending_deallocate_region_ids[1],
            RegionId::new(table_id, 5)
        );
    }

    #[test]
    fn test_convert_plan_allocate_single_region() {
        let group_id = Uuid::new_v4();
        let table_id = 1024;
        let mut next_region_number = 5;
        // 1 source region -> 2 target partition expressions
        let source_regions = vec![create_region_descriptor(table_id, 1, "x", 0, 100)];
        let target_partition_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 100)];
        let allocation_plan = AllocationPlanEntry {
            group_id,
            source_regions: source_regions.clone(),
            target_partition_exprs: target_partition_exprs.clone(),
            transition_map: vec![],
        };

        let result = convert_allocation_plan_to_repartition_plan(
            table_id,
            &mut next_region_number,
            &allocation_plan,
        );
        assert_eq!(result.target_regions.len(), 2);
        assert_eq!(result.allocated_region_ids.len(), 1);
        assert_eq!(result.pending_deallocate_region_ids.len(), 0);
        // First target uses source region id
        assert_eq!(
            result.target_regions[0].region_id,
            RegionId::new(table_id, 1)
        );
        assert_eq!(
            result.target_regions[0].partition_expr,
            target_partition_exprs[0]
        );
        // Second target is newly allocated
        assert_eq!(
            result.target_regions[1].region_id,
            RegionId::new(table_id, 5)
        );
        assert_eq!(
            result.target_regions[1].partition_expr,
            target_partition_exprs[1]
        );
        assert_eq!(next_region_number, 6);
    }

    #[test]
    fn test_convert_plan_allocate_default_source_region() {
        let group_id = Uuid::new_v4();
        let table_id = 1024;
        let mut next_region_number = 5;
        let source_regions = vec![SourceRegionDescriptor::Default {
            region_id: RegionId::new(table_id, 1),
        }];
        let target_partition_exprs = vec![range_expr("x", 0, 50), range_expr("x", 50, 100)];
        let allocation_plan = AllocationPlanEntry {
            group_id,
            source_regions: source_regions.clone(),
            target_partition_exprs: target_partition_exprs.clone(),
            transition_map: vec![vec![0, 1]],
        };

        let result = convert_allocation_plan_to_repartition_plan(
            table_id,
            &mut next_region_number,
            &allocation_plan,
        );

        assert_eq!(result.source_regions, source_regions);
        assert_eq!(result.target_regions.len(), 2);
        assert_eq!(
            result.target_regions[0].region_id,
            RegionId::new(table_id, 1)
        );
        assert_eq!(
            result.target_regions[0].partition_expr,
            target_partition_exprs[0]
        );
        assert_eq!(
            result.target_regions[1].region_id,
            RegionId::new(table_id, 5)
        );
        assert_eq!(
            result.target_regions[1].partition_expr,
            target_partition_exprs[1]
        );
        assert_eq!(
            result.allocated_region_ids,
            vec![RegionId::new(table_id, 5)]
        );
        assert!(result.pending_deallocate_region_ids.is_empty());
        assert_eq!(result.transition_map, vec![vec![0, 1]]);
        assert_eq!(next_region_number, 6);
    }

    #[test]
    fn test_convert_plan_deallocate_to_single_region() {
        let group_id = Uuid::new_v4();
        let table_id = 1024;

        // 3 source regions -> 1 target partition expression
        let source_regions = vec![
            create_region_descriptor(table_id, 1, "x", 0, 100),
            create_region_descriptor(table_id, 2, "x", 100, 200),
            create_region_descriptor(table_id, 3, "x", 200, 300),
        ];
        let target_partition_exprs = vec![range_expr("x", 0, 300)];
        let allocation_plan = AllocationPlanEntry {
            group_id,
            source_regions: source_regions.clone(),
            target_partition_exprs: target_partition_exprs.clone(),
            transition_map: vec![],
        };
        let mut next_region_number = 10;
        let result = convert_allocation_plan_to_repartition_plan(
            table_id,
            &mut next_region_number,
            &allocation_plan,
        );
        assert_eq!(result.target_regions.len(), 1);
        assert_eq!(result.allocated_region_ids.len(), 0);
        assert_eq!(result.pending_deallocate_region_ids.len(), 2);

        // Only first source region is kept
        assert_eq!(
            result.target_regions[0].region_id,
            RegionId::new(table_id, 1)
        );
        assert_eq!(
            result.target_regions[0].partition_expr,
            target_partition_exprs[0]
        );

        // Other regions are pending deallocation
        assert_eq!(
            result.pending_deallocate_region_ids[0],
            RegionId::new(table_id, 2)
        );
        assert_eq!(
            result.pending_deallocate_region_ids[1],
            RegionId::new(table_id, 3)
        );
    }
}