fix: correctly update partition key indices during alter table operations (#6494)

* fix: correctly update partition key indices in alter table operations

Signed-off-by: WenyXu <wenymedia@gmail.com>

* test: add sqlness tests

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2025-07-10 16:08:07 +08:00
committed by GitHub
parent f2c7b09825
commit ffe0da0405
7 changed files with 245 additions and 6 deletions

View File

@@ -20,8 +20,8 @@ use api::v1::region::{alter_request, AlterRequest, RegionRequest, RegionRequestH
use api::v1::AlterTableExpr;
use common_catalog::format_full_table_name;
use common_grpc_expr::alter_expr_to_request;
use common_telemetry::debug;
use common_telemetry::tracing_context::TracingContext;
use common_telemetry::{debug, info};
use futures::future;
use snafu::{ensure, ResultExt};
use store_api::metadata::ColumnMetadata;
@@ -304,5 +304,10 @@ fn build_new_table_info(
| AlterKind::DropDefaults { .. } => {}
}
info!(
"Built new table info: {:?} for table {}, table_id: {}",
new_info.meta, table_name, table_id
);
Ok(new_info)
}

View File

@@ -95,6 +95,18 @@ pub enum Error {
location: Location,
},
#[snafu(display(
"Not allowed to remove partition column {} from table {}",
column_name,
table_name
))]
RemovePartitionColumn {
column_name: String,
table_name: String,
#[snafu(implicit)]
location: Location,
},
#[snafu(display(
"Failed to build column descriptor for table: {}, column: {}",
table_name,
@@ -193,6 +205,7 @@ impl ErrorExt for Error {
StatusCode::EngineExecuteQuery
}
Error::RemoveColumnInIndex { .. }
| Error::RemovePartitionColumn { .. }
| Error::BuildColumnDescriptor { .. }
| Error::InvalidAlterRequest { .. } => StatusCode::InvalidArguments,
Error::CastDefaultValue { source, .. } => source.status_code(),

View File

@@ -649,10 +649,19 @@ impl TableMeta {
msg: format!("Table {table_name} cannot add new columns {column_names:?}"),
})?;
let partition_key_indices = self
.partition_key_indices
.iter()
.map(|idx| table_schema.column_name_by_index(*idx))
// This unwrap is safe since we only add new columns.
.map(|name| new_schema.column_index_by_name(name).unwrap())
.collect();
// value_indices would be generated automatically.
let _ = meta_builder
.schema(Arc::new(new_schema))
.primary_key_indices(primary_key_indices);
.primary_key_indices(primary_key_indices)
.partition_key_indices(partition_key_indices);
Ok(meta_builder)
}
@@ -680,6 +689,14 @@ impl TableMeta {
}
);
ensure!(
!self.partition_key_indices.contains(&index),
error::RemovePartitionColumnSnafu {
column_name: *column_name,
table_name,
}
);
if let Some(ts_index) = timestamp_index {
// Not allowed to remove column in timestamp index.
ensure!(
@@ -729,9 +746,18 @@ impl TableMeta {
.map(|name| new_schema.column_index_by_name(name).unwrap())
.collect();
let partition_key_indices = self
.partition_key_indices
.iter()
.map(|idx| table_schema.column_name_by_index(*idx))
// This unwrap is safe since we don't allow removing a partition key column.
.map(|name| new_schema.column_index_by_name(name).unwrap())
.collect();
let _ = meta_builder
.schema(Arc::new(new_schema))
.primary_key_indices(primary_key_indices);
.primary_key_indices(primary_key_indices)
.partition_key_indices(partition_key_indices);
Ok(meta_builder)
}
@@ -1334,6 +1360,8 @@ fn unset_column_skipping_index_options(
#[cfg(test)]
mod tests {
use std::assert_matches::assert_matches;
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use datatypes::data_type::ConcreteDataType;
@@ -1342,6 +1370,7 @@ mod tests {
};
use super::*;
use crate::Error;
/// Create a test schema with 3 columns: `[col1 int32, ts timestampmills, col2 int32]`.
fn new_test_schema() -> Schema {
@@ -1419,6 +1448,11 @@ mod tests {
ConcreteDataType::string_datatype(),
true,
);
let yet_another_field = ColumnSchema::new(
"yet_another_field_after_ts",
ConcreteDataType::int64_datatype(),
true,
);
let alter_kind = AlterKind::AddColumns {
columns: vec![
AddColumnRequest {
@@ -1435,6 +1469,14 @@ mod tests {
}),
add_if_not_exists: false,
},
AddColumnRequest {
column_schema: yet_another_field,
is_key: true,
location: Some(AddColumnLocation::After {
column_name: "ts".to_string(),
}),
add_if_not_exists: false,
},
],
};
@@ -1790,6 +1832,29 @@ mod tests {
assert_eq!(StatusCode::InvalidArguments, err.status_code());
}
#[test]
fn test_remove_partition_column() {
let schema = Arc::new(new_test_schema());
let meta = TableMetaBuilder::empty()
.schema(schema)
.primary_key_indices(vec![])
.partition_key_indices(vec![0])
.engine("engine")
.next_column_id(3)
.build()
.unwrap();
// Remove column in primary key.
let alter_kind = AlterKind::DropColumns {
names: vec![String::from("col1")],
};
let err = meta
.builder_with_alter_kind("my_table", &alter_kind)
.err()
.unwrap();
assert_matches!(err, Error::RemovePartitionColumn { .. });
}
#[test]
fn test_change_key_column_data_type() {
let schema = Arc::new(new_test_schema());
@@ -1855,6 +1920,8 @@ mod tests {
let meta = TableMetaBuilder::empty()
.schema(schema)
.primary_key_indices(vec![0])
// partition col: col1, col2
.partition_key_indices(vec![0, 2])
.engine("engine")
.next_column_id(3)
.build()
@@ -1870,11 +1937,19 @@ mod tests {
.map(|column_schema| column_schema.name.clone())
.collect();
assert_eq!(
&["my_tag_first", "col1", "ts", "my_field_after_ts", "col2"],
&[
"my_tag_first", // primary key column
"col1", // partition column
"ts", // timestamp column
"yet_another_field_after_ts", // primary key column
"my_field_after_ts", // value column
"col2", // partition column
],
&names[..]
);
assert_eq!(&[0, 1], &new_meta.primary_key_indices[..]);
assert_eq!(&[2, 3, 4], &new_meta.value_indices[..]);
assert_eq!(&[0, 1, 3], &new_meta.primary_key_indices[..]);
assert_eq!(&[2, 4, 5], &new_meta.value_indices[..]);
assert_eq!(&[1, 5], &new_meta.partition_key_indices[..]);
}
#[test]

View File

@@ -174,3 +174,80 @@ DROP TABLE t;
Affected Rows: 0
CREATE TABLE my_table (
a INT PRIMARY KEY,
b STRING,
ts TIMESTAMP TIME INDEX,
)
PARTITION ON COLUMNS (a) (
a < 1000,
a >= 1000 AND a < 2000,
a >= 2000
);
Affected Rows: 0
INSERT INTO my_table VALUES
(100, 'a', 1),
(200, 'b', 2),
(1100, 'c', 3),
(1200, 'd', 4),
(2000, 'e', 5),
(2100, 'f', 6),
(2200, 'g', 7),
(2400, 'h', 8);
Affected Rows: 8
SELECT * FROM my_table WHERE a > 100 order by a;
+------+---+-------------------------+
| a | b | ts |
+------+---+-------------------------+
| 200 | b | 1970-01-01T00:00:00.002 |
| 1100 | c | 1970-01-01T00:00:00.003 |
| 1200 | d | 1970-01-01T00:00:00.004 |
| 2000 | e | 1970-01-01T00:00:00.005 |
| 2100 | f | 1970-01-01T00:00:00.006 |
| 2200 | g | 1970-01-01T00:00:00.007 |
| 2400 | h | 1970-01-01T00:00:00.008 |
+------+---+-------------------------+
SELECT count(*) FROM my_table WHERE a > 100;
+----------+
| count(*) |
+----------+
| 7 |
+----------+
ALTER TABLE my_table ADD COLUMN c STRING FIRST;
Affected Rows: 0
SELECT * FROM my_table WHERE a > 100 order by a;
+---+------+---+-------------------------+
| c | a | b | ts |
+---+------+---+-------------------------+
| | 200 | b | 1970-01-01T00:00:00.002 |
| | 1100 | c | 1970-01-01T00:00:00.003 |
| | 1200 | d | 1970-01-01T00:00:00.004 |
| | 2000 | e | 1970-01-01T00:00:00.005 |
| | 2100 | f | 1970-01-01T00:00:00.006 |
| | 2200 | g | 1970-01-01T00:00:00.007 |
| | 2400 | h | 1970-01-01T00:00:00.008 |
+---+------+---+-------------------------+
SELECT count(*) FROM my_table WHERE a > 100;
+----------+
| count(*) |
+----------+
| 7 |
+----------+
DROP TABLE my_table;
Affected Rows: 0

View File

@@ -47,3 +47,36 @@ SELECT * FROM t;
ALTER TABLE t ADD COLUMN x int xxx;
DROP TABLE t;
CREATE TABLE my_table (
a INT PRIMARY KEY,
b STRING,
ts TIMESTAMP TIME INDEX,
)
PARTITION ON COLUMNS (a) (
a < 1000,
a >= 1000 AND a < 2000,
a >= 2000
);
INSERT INTO my_table VALUES
(100, 'a', 1),
(200, 'b', 2),
(1100, 'c', 3),
(1200, 'd', 4),
(2000, 'e', 5),
(2100, 'f', 6),
(2200, 'g', 7),
(2400, 'h', 8);
SELECT * FROM my_table WHERE a > 100 order by a;
SELECT count(*) FROM my_table WHERE a > 100;
ALTER TABLE my_table ADD COLUMN c STRING FIRST;
SELECT * FROM my_table WHERE a > 100 order by a;
SELECT count(*) FROM my_table WHERE a > 100;
DROP TABLE my_table;

View File

@@ -31,3 +31,24 @@ DROP TABLE test;
Affected Rows: 0
CREATE TABLE my_table (
a INT PRIMARY KEY,
b STRING,
ts TIMESTAMP TIME INDEX,
)
PARTITION ON COLUMNS (a) (
a < 1000,
a >= 1000 AND a < 2000,
a >= 2000
);
Affected Rows: 0
ALTER TABLE my_table DROP COLUMN a;
Error: 1004(InvalidArguments), Not allowed to remove index column a from table my_table
DROP TABLE my_table;
Affected Rows: 0

View File

@@ -11,3 +11,18 @@ SELECT * FROM test;
ALTER TABLE test DROP COLUMN j;
DROP TABLE test;
CREATE TABLE my_table (
a INT PRIMARY KEY,
b STRING,
ts TIMESTAMP TIME INDEX,
)
PARTITION ON COLUMNS (a) (
a < 1000,
a >= 1000 AND a < 2000,
a >= 2000
);
ALTER TABLE my_table DROP COLUMN a;
DROP TABLE my_table;