Skip to content

Commit fa8c948

Browse files
authored
feat: Add RemovePartitionSpecs table update (#804)
* RemovePartitionSpecs update * address comments * fix comments
1 parent f90dbff commit fa8c948

File tree

2 files changed

+76
-0
lines changed

2 files changed

+76
-0
lines changed

crates/iceberg/src/catalog/mod.rs

+24
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,12 @@ pub enum TableUpdate {
440440
/// Properties to remove
441441
removals: Vec<String>,
442442
},
443+
/// Remove partition specs
444+
#[serde(rename_all = "kebab-case")]
445+
RemovePartitionSpecs {
446+
/// Partition spec ids to remove.
447+
spec_ids: Vec<i32>,
448+
},
443449
/// Set statistics for a snapshot
444450
#[serde(with = "_serde_set_statistics")]
445451
SetStatistics {
@@ -494,6 +500,9 @@ impl TableUpdate {
494500
TableUpdate::UpgradeFormatVersion { format_version } => {
495501
builder.upgrade_format_version(format_version)
496502
}
503+
TableUpdate::RemovePartitionSpecs { spec_ids } => {
504+
builder.remove_partition_specs(&spec_ids)
505+
}
497506
TableUpdate::SetStatistics { statistics } => Ok(builder.set_statistics(statistics)),
498507
TableUpdate::RemoveStatistics { snapshot_id } => {
499508
Ok(builder.remove_statistics(snapshot_id))
@@ -1910,6 +1919,21 @@ mod tests {
19101919
);
19111920
}
19121921

1922+
#[test]
1923+
fn test_remove_partition_specs_update() {
1924+
test_serde_json(
1925+
r#"
1926+
{
1927+
"action": "remove-partition-specs",
1928+
"spec-ids": [1, 2]
1929+
}
1930+
"#,
1931+
TableUpdate::RemovePartitionSpecs {
1932+
spec_ids: vec![1, 2],
1933+
},
1934+
);
1935+
}
1936+
19131937
#[test]
19141938
fn test_set_statistics_file() {
19151939
test_serde_json(

crates/iceberg/src/spec/table_metadata_builder.rs

+52
Original file line numberDiff line numberDiff line change
@@ -784,6 +784,36 @@ impl TableMetadataBuilder {
784784
.set_default_partition_spec(Self::LAST_ADDED)
785785
}
786786

787+
/// Remove partition specs by their ids from the table metadata.
788+
/// Does nothing if a spec id is not present. Active partition specs
789+
/// should not be removed.
790+
///
791+
/// # Errors
792+
/// - Cannot remove the default partition spec.
793+
pub fn remove_partition_specs(mut self, spec_ids: &[i32]) -> Result<Self> {
794+
if spec_ids.contains(&self.metadata.default_spec.spec_id()) {
795+
return Err(Error::new(
796+
ErrorKind::DataInvalid,
797+
"Cannot remove default partition spec",
798+
));
799+
}
800+
801+
let mut removed_specs = Vec::with_capacity(spec_ids.len());
802+
spec_ids.iter().for_each(|id| {
803+
if self.metadata.partition_specs.remove(id).is_some() {
804+
removed_specs.push(*id);
805+
}
806+
});
807+
808+
if !removed_specs.is_empty() {
809+
self.changes.push(TableUpdate::RemovePartitionSpecs {
810+
spec_ids: removed_specs,
811+
});
812+
}
813+
814+
Ok(self)
815+
}
816+
787817
/// Add a sort order to the table metadata.
788818
///
789819
/// The spec is bound eagerly to the current schema and must be valid for it.
@@ -1584,6 +1614,21 @@ mod tests {
15841614
pretty_assertions::assert_eq!(build_result.changes[0], TableUpdate::AddSpec {
15851615
spec: expected_change
15861616
});
1617+
1618+
// Remove the spec
1619+
let build_result = build_result
1620+
.metadata
1621+
.into_builder(Some(
1622+
"s3://bucket/test/location/metadata/metadata1.json".to_string(),
1623+
))
1624+
.remove_partition_specs(&[1])
1625+
.unwrap()
1626+
.build()
1627+
.unwrap();
1628+
1629+
assert_eq!(build_result.changes.len(), 1);
1630+
assert_eq!(build_result.metadata.partition_specs.len(), 1);
1631+
assert!(build_result.metadata.partition_spec_by_id(1).is_none());
15871632
}
15881633

15891634
#[test]
@@ -2193,6 +2238,13 @@ mod tests {
21932238
.contains("Cannot add snapshot with sequence number"));
21942239
}
21952240

2241+
#[test]
2242+
fn test_default_spec_cannot_be_removed() {
2243+
let builder = builder_without_changes(FormatVersion::V2);
2244+
2245+
builder.remove_partition_specs(&[0]).unwrap_err();
2246+
}
2247+
21962248
#[test]
21972249
fn test_statistics() {
21982250
let builder = builder_without_changes(FormatVersion::V2);

0 commit comments

Comments
 (0)