diff --git a/arrow-schema/src/fields.rs b/arrow-schema/src/fields.rs index 9cb56b95dda3..70cb1968e9a4 100644 --- a/arrow-schema/src/fields.rs +++ b/arrow-schema/src/fields.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::{ArrowError, Field, FieldRef}; +use crate::{ArrowError, Field, FieldRef, SchemaBuilder}; use std::ops::Deref; use std::sync::Arc; @@ -100,6 +100,12 @@ impl Fields { } /// Remove a field by index and return it. + /// + /// # Panic + /// + /// Panics if `index` is out of bounds. + /// + /// # Example /// ``` /// use arrow_schema::{DataType, Field, Fields}; /// let mut fields = Fields::from(vec![ @@ -108,18 +114,14 @@ impl Fields { /// Field::new("c", DataType::Utf8, false), /// ]); /// assert_eq!(fields.len(), 3); - /// assert_eq!(fields.remove(1).unwrap(), Field::new("b", DataType::Int8, false).into()); + /// assert_eq!(fields.remove(1), Field::new("b", DataType::Int8, false).into()); /// assert_eq!(fields.len(), 2); /// ``` - pub fn remove(&mut self, index: usize) -> Option { - if index >= self.len() { - return None; - } - - let mut new_fields = self.0.iter().cloned().collect::>(); - let field = new_fields.remove(index); - self.0 = Arc::from(new_fields); - Some(field) + pub fn remove(&mut self, index: usize) -> FieldRef { + let mut builder = SchemaBuilder::from(Fields::from(&*self.0)); + let field = builder.remove(index); + *self = builder.finish().fields; + field } } diff --git a/arrow-schema/src/schema.rs b/arrow-schema/src/schema.rs index c06498961b2a..e30256524868 100644 --- a/arrow-schema/src/schema.rs +++ b/arrow-schema/src/schema.rs @@ -336,6 +336,12 @@ impl Schema { /// Remove field by name and return it. Recommend to use [`SchemaBuilder`] /// if you are looking to remove multiple columns, as this will save allocations. /// + /// # Panic + /// + /// Panics if `index` is out of bounds. + /// + /// # Example + /// /// ``` /// use arrow_schema::{DataType, Field, Schema}; /// let mut schema = Schema::new(vec![ @@ -344,15 +350,11 @@ impl Schema { /// Field::new("c", DataType::Utf8, false), /// ]); /// assert_eq!(schema.fields.len(), 3); - /// assert_eq!(schema.remove_field_with_name("b").unwrap(), Field::new("b", DataType::Int8, false).into()); + /// assert_eq!(schema.remove(1), Field::new("b", DataType::Int8, false).into()); /// assert_eq!(schema.fields.len(), 2); /// ``` - pub fn remove_field_with_name(&mut self, name: &str) -> Option { - if let Some((idx, _)) = self.fields.find(name) { - self.fields.remove(idx) - } else { - None - } + pub fn remove(&mut self, index: usize) -> FieldRef { + self.fields.remove(index) } }