JeonDaehong commented on code in PR #3441: URL: https://github.com/apache/iceberg-rust/pull/3441#discussion_r4348335151
########## crates/iceberg/src/transaction/update_partition_spec.rs: ########## @@ -0,0 +1,1267 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::error::{Error, ErrorKind, Result, invalid_data}; +use crate::spec::{ + FormatVersion, PartitionField, PartitionSpec, SchemaRef, Transform, UnboundPartitionSpec, +}; +use crate::table::Table; +use crate::transaction::{ActionCommit, TransactionAction}; +use crate::{TableRequirement, TableUpdate}; + +enum PendingChange { + Add { + name: Option<String>, + source_name: String, + transform: Transform, + }, + RemoveByName(String), + RemoveByTransform { + source_name: String, + transform: Transform, + }, + Rename { + name: String, + new_name: String, + }, +} + +/// Transaction action that evolves the table's partition spec. +/// +/// Changes are validated and resolved against the table metadata at commit time, so a commit +/// that is retried after a concurrent update re-applies them to the refreshed spec. A new field +/// with the same source and transform as a field of an earlier spec, and the same name if one +/// is given, reuses that field and its field id. Format version 1 tables are not supported yet. +/// +/// # Example +/// +/// ```ignore +/// let tx = Transaction::new(&table); +/// let action = tx +/// .update_partition_spec() +/// .add_field("id", Transform::Bucket(16)) +/// .remove_field("category") +/// .rename_field("ts_day", "day"); +/// let tx = action.apply(tx).unwrap(); +/// let table = tx.commit(&catalog).await.unwrap(); +/// ``` +pub struct UpdatePartitionSpecAction { + changes: Vec<PendingChange>, + case_sensitive: bool, + set_as_default: bool, +} + +impl UpdatePartitionSpecAction { + pub(crate) fn new() -> Self { + Self { + changes: Vec::new(), + case_sensitive: true, + set_as_default: true, + } + } + + /// Partition by the value of a source column, in a field named after the column. + pub fn add_identity(self, source_name: impl Into<String>) -> Self { + self.add_field(source_name, Transform::Identity) + } + + /// Partition by `transform` applied to a source column. + /// + /// The field takes the name of an equivalent field from an earlier spec, if there is one. + /// Otherwise it gets the default name for the transform, such as `id_bucket_16` or `ts_day`. + pub fn add_field(mut self, source_name: impl Into<String>, transform: Transform) -> Self { + self.changes.push(PendingChange::Add { + name: None, + source_name: source_name.into(), + transform, + }); + self + } + + /// Partition by `transform` applied to a source column, in a field with the given name. + pub fn add_field_with_name( + mut self, + name: impl Into<String>, + source_name: impl Into<String>, + transform: Transform, + ) -> Self { + self.changes.push(PendingChange::Add { + name: Some(name.into()), + source_name: source_name.into(), + transform, + }); + self + } + + /// Remove the partition field with the given name. + pub fn remove_field(mut self, name: impl Into<String>) -> Self { + self.changes.push(PendingChange::RemoveByName(name.into())); + self + } + + /// Remove the partition field that applies `transform` to a source column. + pub fn remove_field_by_transform( + mut self, + source_name: impl Into<String>, + transform: Transform, + ) -> Self { + self.changes.push(PendingChange::RemoveByTransform { + source_name: source_name.into(), + transform, + }); + self + } + + /// Rename a partition field of the current spec. + pub fn rename_field(mut self, name: impl Into<String>, new_name: impl Into<String>) -> Self { + self.changes.push(PendingChange::Rename { + name: name.into(), + new_name: new_name.into(), + }); + self + } + + /// Set whether source column names match case-sensitively. Defaults to `true`. + pub fn case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Add the new spec to the table without making it the default spec. + pub fn add_non_default_spec(mut self) -> Self { + self.set_as_default = false; + self + } + + fn resolve_spec<'a>( + &self, + schema: &'a SchemaRef, + spec: &'a PartitionSpec, + historical_specs: impl IntoIterator<Item = &'a PartitionSpec>, + last_assigned_partition_id: i32, + ) -> Result<UnboundPartitionSpec> { + let mut update = PartitionSpecUpdate::new( + schema, + spec, + historical_specs, + last_assigned_partition_id, + self.case_sensitive, + )?; + for change in &self.changes { + match change { + PendingChange::Add { + name, + source_name, + transform, + } => update.add_field(name.as_deref(), source_name, *transform)?, + PendingChange::RemoveByName(name) => update.remove_field(name)?, + PendingChange::RemoveByTransform { + source_name, + transform, + } => update.remove_field_by_transform(source_name, *transform)?, + PendingChange::Rename { name, new_name } => update.rename_field(name, new_name)?, + } + } + update.apply() + } +} + +#[async_trait] +impl TransactionAction for UpdatePartitionSpecAction { + async fn commit(self: Arc<Self>, table: &Table) -> Result<ActionCommit> { + let metadata = table.metadata(); + if metadata.format_version() == FormatVersion::V1 { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "Updating the partition spec of a format version 1 table is not supported", + )); + } + + let spec = self.resolve_spec( + metadata.current_schema(), + metadata.default_partition_spec(), + metadata.partition_specs_iter().map(|spec| spec.as_ref()), + metadata.last_partition_id(), + )?; + + let mut updates = vec![TableUpdate::AddSpec { spec }]; + let mut requirements = vec![TableRequirement::LastAssignedPartitionIdMatch { + last_assigned_partition_id: metadata.last_partition_id(), + }]; + if self.set_as_default { + updates.push(TableUpdate::SetDefaultSpec { spec_id: -1 }); + requirements.push(TableRequirement::DefaultSpecIdMatch { + default_spec_id: metadata.default_partition_spec_id(), + }); + } + + Ok(ActionCommit::new(updates, requirements)) + } +} + +/// Applies partition field changes to a base spec, following Java's `BaseUpdatePartitionSpec` +/// so that both produce the same spec for the same changes. +struct PartitionSpecUpdate<'a> { + schema: &'a SchemaRef, + spec: &'a PartitionSpec, + historical_fields: Vec<&'a PartitionField>, + case_sensitive: bool, + last_assigned_partition_id: i32, + name_to_field: HashMap<&'a str, &'a PartitionField>, + transform_to_field: HashMap<(i32, Transform), &'a PartitionField>, + adds: Vec<PartitionField>, + added_names: HashSet<String>, + transform_to_added_name: HashMap<(i32, Transform), String>, + added_time_field_names: HashMap<i32, String>, + deletes: HashSet<i32>, + renames: HashMap<String, String>, +} + +impl<'a> PartitionSpecUpdate<'a> { + fn new( + schema: &'a SchemaRef, + spec: &'a PartitionSpec, + historical_specs: impl IntoIterator<Item = &'a PartitionSpec>, + last_assigned_partition_id: i32, + case_sensitive: bool, + ) -> Result<Self> { + if let Some(field) = spec + .fields() + .iter() + .find(|field| field.transform == Transform::Unknown) + { + return Err(invalid_data!( + "Cannot update partition spec with unknown transform: {}", + field.name + )); + } + + let mut historical_specs: Vec<_> = historical_specs.into_iter().collect(); + historical_specs.sort_by_key(|spec| spec.spec_id()); + + Ok(Self { + schema, + spec, + historical_fields: historical_specs + .into_iter() + .flat_map(|spec| spec.fields()) + .collect(), + case_sensitive, + last_assigned_partition_id, + name_to_field: spec + .fields() + .iter() + .map(|field| (field.name.as_str(), field)) + .collect(), + transform_to_field: spec + .fields() + .iter() + .map(|field| ((field.source_id, field.transform), field)) + .collect(), + adds: Vec::new(), + added_names: HashSet::new(), + transform_to_added_name: HashMap::new(), + added_time_field_names: HashMap::new(), + deletes: HashSet::new(), + renames: HashMap::new(), + }) + } + + fn add_field( + &mut self, + name: Option<&str>, + source_name: &str, + transform: Transform, + ) -> Result<()> { + if let Some(name) = name + && self.added_names.contains(name) + { + return Err(invalid_data!( + "Cannot add duplicate partition field: {name}" + )); + } + if matches!(transform, Transform::Void | Transform::Unknown) { + return Err(invalid_data!( + "Cannot add partition field with transform {transform}" + )); + } + + let source_id = self.source_id(source_name)?; + let key = (source_id, transform); + if let Some(existing) = self.transform_to_field.get(&key).copied() { + if self.deletes.remove(&existing.field_id) { Review Comment: This restores the removed field for every transform, which is what Java intends (`rewriteDeleteAndAddField`). Java only gets there for some of them, though. Its check is `existing.transform().equals(...)`, and for time transforms the bound `day` and the new `day` aren't equal, so Java falls through to the duplicate error: | same update | Java (1.12.0-SNAPSHOT) | here | |---|---|---| | remove `category`, add `identity(category)` as `c` | `1000:c` | same | | remove `shard`, add `bucket[16](id)` as `b` | `1002:b` | same | | remove `ts_day`, add `day(ts)` | **error**: duplicate, conflicts with `1001: ts_day` | `1001:ts_day` | | remove `ts_day`, add `day(ts)` as `d` | **error** | `1001:d` | I think the behavior here is the better one. But since the tests are a port of Java's, nothing covers the time-transform case. A test for it would make the difference from Java deliberate rather than accidental. ########## crates/iceberg/src/transaction/update_partition_spec.rs: ########## @@ -0,0 +1,1267 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::error::{Error, ErrorKind, Result, invalid_data}; +use crate::spec::{ + FormatVersion, PartitionField, PartitionSpec, SchemaRef, Transform, UnboundPartitionSpec, +}; +use crate::table::Table; +use crate::transaction::{ActionCommit, TransactionAction}; +use crate::{TableRequirement, TableUpdate}; + +enum PendingChange { + Add { + name: Option<String>, + source_name: String, + transform: Transform, + }, + RemoveByName(String), + RemoveByTransform { + source_name: String, + transform: Transform, + }, + Rename { + name: String, + new_name: String, + }, +} + +/// Transaction action that evolves the table's partition spec. +/// +/// Changes are validated and resolved against the table metadata at commit time, so a commit +/// that is retried after a concurrent update re-applies them to the refreshed spec. A new field +/// with the same source and transform as a field of an earlier spec, and the same name if one +/// is given, reuses that field and its field id. Format version 1 tables are not supported yet. +/// +/// # Example +/// +/// ```ignore +/// let tx = Transaction::new(&table); +/// let action = tx +/// .update_partition_spec() +/// .add_field("id", Transform::Bucket(16)) +/// .remove_field("category") +/// .rename_field("ts_day", "day"); +/// let tx = action.apply(tx).unwrap(); +/// let table = tx.commit(&catalog).await.unwrap(); +/// ``` +pub struct UpdatePartitionSpecAction { + changes: Vec<PendingChange>, + case_sensitive: bool, + set_as_default: bool, +} + +impl UpdatePartitionSpecAction { + pub(crate) fn new() -> Self { + Self { + changes: Vec::new(), + case_sensitive: true, + set_as_default: true, + } + } + + /// Partition by the value of a source column, in a field named after the column. + pub fn add_identity(self, source_name: impl Into<String>) -> Self { + self.add_field(source_name, Transform::Identity) + } + + /// Partition by `transform` applied to a source column. + /// + /// The field takes the name of an equivalent field from an earlier spec, if there is one. + /// Otherwise it gets the default name for the transform, such as `id_bucket_16` or `ts_day`. + pub fn add_field(mut self, source_name: impl Into<String>, transform: Transform) -> Self { + self.changes.push(PendingChange::Add { + name: None, + source_name: source_name.into(), + transform, + }); + self + } + + /// Partition by `transform` applied to a source column, in a field with the given name. + pub fn add_field_with_name( + mut self, + name: impl Into<String>, + source_name: impl Into<String>, + transform: Transform, + ) -> Self { + self.changes.push(PendingChange::Add { + name: Some(name.into()), + source_name: source_name.into(), + transform, + }); + self + } + + /// Remove the partition field with the given name. + pub fn remove_field(mut self, name: impl Into<String>) -> Self { + self.changes.push(PendingChange::RemoveByName(name.into())); + self + } + + /// Remove the partition field that applies `transform` to a source column. + pub fn remove_field_by_transform( + mut self, + source_name: impl Into<String>, + transform: Transform, + ) -> Self { + self.changes.push(PendingChange::RemoveByTransform { + source_name: source_name.into(), + transform, + }); + self + } + + /// Rename a partition field of the current spec. + pub fn rename_field(mut self, name: impl Into<String>, new_name: impl Into<String>) -> Self { + self.changes.push(PendingChange::Rename { + name: name.into(), + new_name: new_name.into(), + }); + self + } + + /// Set whether source column names match case-sensitively. Defaults to `true`. + pub fn case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Add the new spec to the table without making it the default spec. + pub fn add_non_default_spec(mut self) -> Self { + self.set_as_default = false; + self + } + + fn resolve_spec<'a>( + &self, + schema: &'a SchemaRef, + spec: &'a PartitionSpec, + historical_specs: impl IntoIterator<Item = &'a PartitionSpec>, + last_assigned_partition_id: i32, + ) -> Result<UnboundPartitionSpec> { + let mut update = PartitionSpecUpdate::new( + schema, + spec, + historical_specs, + last_assigned_partition_id, + self.case_sensitive, + )?; + for change in &self.changes { + match change { + PendingChange::Add { + name, + source_name, + transform, + } => update.add_field(name.as_deref(), source_name, *transform)?, + PendingChange::RemoveByName(name) => update.remove_field(name)?, + PendingChange::RemoveByTransform { + source_name, + transform, + } => update.remove_field_by_transform(source_name, *transform)?, + PendingChange::Rename { name, new_name } => update.rename_field(name, new_name)?, + } + } + update.apply() + } +} + +#[async_trait] +impl TransactionAction for UpdatePartitionSpecAction { + async fn commit(self: Arc<Self>, table: &Table) -> Result<ActionCommit> { + let metadata = table.metadata(); + if metadata.format_version() == FormatVersion::V1 { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "Updating the partition spec of a format version 1 table is not supported", + )); + } + + let spec = self.resolve_spec( + metadata.current_schema(), + metadata.default_partition_spec(), + metadata.partition_specs_iter().map(|spec| spec.as_ref()), + metadata.last_partition_id(), + )?; + + let mut updates = vec![TableUpdate::AddSpec { spec }]; + let mut requirements = vec![TableRequirement::LastAssignedPartitionIdMatch { + last_assigned_partition_id: metadata.last_partition_id(), + }]; + if self.set_as_default { + updates.push(TableUpdate::SetDefaultSpec { spec_id: -1 }); + requirements.push(TableRequirement::DefaultSpecIdMatch { + default_spec_id: metadata.default_partition_spec_id(), + }); + } + + Ok(ActionCommit::new(updates, requirements)) + } +} + +/// Applies partition field changes to a base spec, following Java's `BaseUpdatePartitionSpec` +/// so that both produce the same spec for the same changes. +struct PartitionSpecUpdate<'a> { + schema: &'a SchemaRef, + spec: &'a PartitionSpec, + historical_fields: Vec<&'a PartitionField>, + case_sensitive: bool, + last_assigned_partition_id: i32, + name_to_field: HashMap<&'a str, &'a PartitionField>, + transform_to_field: HashMap<(i32, Transform), &'a PartitionField>, + adds: Vec<PartitionField>, + added_names: HashSet<String>, + transform_to_added_name: HashMap<(i32, Transform), String>, + added_time_field_names: HashMap<i32, String>, + deletes: HashSet<i32>, + renames: HashMap<String, String>, +} + +impl<'a> PartitionSpecUpdate<'a> { + fn new( + schema: &'a SchemaRef, + spec: &'a PartitionSpec, + historical_specs: impl IntoIterator<Item = &'a PartitionSpec>, + last_assigned_partition_id: i32, + case_sensitive: bool, + ) -> Result<Self> { + if let Some(field) = spec + .fields() + .iter() + .find(|field| field.transform == Transform::Unknown) + { + return Err(invalid_data!( + "Cannot update partition spec with unknown transform: {}", + field.name + )); + } + + let mut historical_specs: Vec<_> = historical_specs.into_iter().collect(); + historical_specs.sort_by_key(|spec| spec.spec_id()); + + Ok(Self { + schema, + spec, + historical_fields: historical_specs + .into_iter() + .flat_map(|spec| spec.fields()) + .collect(), + case_sensitive, + last_assigned_partition_id, + name_to_field: spec + .fields() + .iter() + .map(|field| (field.name.as_str(), field)) + .collect(), + transform_to_field: spec + .fields() + .iter() + .map(|field| ((field.source_id, field.transform), field)) + .collect(), + adds: Vec::new(), + added_names: HashSet::new(), + transform_to_added_name: HashMap::new(), + added_time_field_names: HashMap::new(), + deletes: HashSet::new(), + renames: HashMap::new(), + }) + } + + fn add_field( + &mut self, + name: Option<&str>, + source_name: &str, + transform: Transform, + ) -> Result<()> { + if let Some(name) = name + && self.added_names.contains(name) + { + return Err(invalid_data!( + "Cannot add duplicate partition field: {name}" + )); + } + if matches!(transform, Transform::Void | Transform::Unknown) { + return Err(invalid_data!( + "Cannot add partition field with transform {transform}" + )); + } + + let source_id = self.source_id(source_name)?; + let key = (source_id, transform); + if let Some(existing) = self.transform_to_field.get(&key).copied() { + if self.deletes.remove(&existing.field_id) { + // Re-adding a field removed in this update restores it, renamed if asked. + return match name { + Some(name) if name != existing.name => self.rename_field(&existing.name, name), + _ => Ok(()), + }; + } + return Err(invalid_data!( + "Cannot add duplicate partition field {transform}({source_name}), conflicts with {}", Review Comment: The description says the error messages match Java, but the duplicate ones are worded differently. Java reports `Cannot add duplicate partition field d=day(ref(name="ts")), conflicts with 1001: ts_day: day(2)`, with the requested name and the existing field's id. Here it's `Cannot add duplicate partition field day(ts), conflicts with ts_day`. That's not a problem in itself, just the description being slightly off. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
