Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Get child array by column name
Browse files Browse the repository at this point in the history
  • Loading branch information
yjshen committed Sep 17, 2021
1 parent 06c45f8 commit cddb512
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 1 deletion.
28 changes: 28 additions & 0 deletions src/array/struct_.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 {
Expand Down
6 changes: 5 additions & 1 deletion tests/it/array/growable/struct_.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down

0 comments on commit cddb512

Please sign in to comment.