diff --git a/src/array/struct_.rs b/src/array/struct_.rs index ddb18724135..f5964908f8a 100644 --- a/src/array/struct_.rs +++ b/src/array/struct_.rs @@ -8,6 +8,7 @@ use crate::{ }; use super::{ffi::ToFfi, new_empty_array, new_null_array, Array, FromFfi}; +use crate::array::ArrayRef; /// A [`StructArray`] is a nested [`Array`] with an optional validity representing /// multiple [`Array`] with the same number of rows. @@ -134,10 +135,37 @@ impl StructArray { &self.values } + /// Returns the field at `pos`. + pub fn value(&self, pos: usize) -> &ArrayRef { + &self.values[pos] + } + + /// Return the number of fields in this struct array + pub fn num_columns(&self) -> usize { + self.values.len() + } + /// Returns the fields of this [`StructArray`]. pub fn fields(&self) -> &[Field] { Self::get_fields(&self.data_type) } + + /// Return field names in this struct array + pub fn column_names(&self) -> Vec<&str> { + self.fields().iter().map(|f| f.name.as_str()).collect() + } + + /// Return child array whose field name equals to column_name + /// + /// Note: A schema can currently have duplicate field names, in which case + /// the first field will always be selected. + /// This issue will be addressed in [ARROW-11178](https://issues.apache.org/jira/browse/ARROW-11178) + pub fn column_by_name(&self, column_name: &str) -> Option<&ArrayRef> { + self.column_names() + .iter() + .position(|c| c == &column_name) + .map(|pos| self.value(pos)) + } } impl StructArray { diff --git a/tests/it/array/growable/struct_.rs b/tests/it/array/growable/struct_.rs index 0dd85e0dcd1..5b7cdc47972 100644 --- a/tests/it/array/growable/struct_.rs +++ b/tests/it/array/growable/struct_.rs @@ -45,7 +45,11 @@ fn basic() { vec![values[0].slice(1, 2).into(), values[1].slice(1, 2).into()], None, ); - assert_eq!(result, expected) + assert_eq!(result, expected); + assert_eq!(result.column_names(), expected.column_names()); + assert_eq!(result.num_columns(), expected.num_columns()); + assert_eq!(result.column_by_name("f1"), expected.column_by_name("f1")); + assert_eq!(result.column_by_name("f2"), expected.column_by_name("f2")); } #[test]