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>
Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
Weny Xu
2025-07-10 16:08:07 +08:00
committed by Yingwen
parent 92fa33c250
commit ecacf1333e
6 changed files with 239 additions and 5 deletions

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

@@ -645,10 +645,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)
}
@@ -676,6 +685,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!(
@@ -725,9 +742,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)
}
@@ -1300,6 +1326,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;
@@ -1308,6 +1336,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 {
@@ -1385,6 +1414,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 {
@@ -1401,6 +1435,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,
},
],
};
@@ -1756,6 +1798,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());
@@ -1821,6 +1886,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()
@@ -1836,11 +1903,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;