feat: column_name based index access for RecordBatch and StructArray

arrow-array/src/array/struct_array.rs

@@ -20,7 +20,7 @@ use arrow_buffer::buffer::buffer_bin_or;
 use arrow_buffer::Buffer;
 use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType, Field};
-use std::any::Any;
+use std::{any::Any, ops::Index};
 
 /// A nested array type where each child (called *field*) is represented by a separate
 /// array.
@@ -296,6 +296,23 @@ impl From<(Vec<(Field, ArrayRef)>, Buffer)> for StructArray {
     }
 }
 
+impl Index<&str> for StructArray {
+    type Output = ArrayRef;
+
+    /// Get a reference to a column's array by 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)
+    ///
+    /// # Panics
+    ///
+    /// Panics if the name is not in the schema.
+    fn index(&self, name: &str) -> &Self::Output {
+        self.column_by_name(name).unwrap()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -352,6 +369,26 @@ mod tests {
         assert_eq!(0, struct_array.offset());
     }
 
+    /// validates that struct can be accessed using `column_name` as index i.e. `struct_array["column_name"]`.
+    #[test]
+    fn test_struct_array_index_access() {
+        let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
+        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
+
+        let struct_array = StructArray::from(vec![
+            (
+                Field::new("b", DataType::Boolean, false),
+                boolean.clone() as ArrayRef,
+            ),
+            (
+                Field::new("c", DataType::Int32, false),
+                int.clone() as ArrayRef,
+            ),
+        ]);
+        assert_eq!(struct_array["b"].as_ref(), boolean.as_ref());
+        assert_eq!(struct_array["c"].as_ref(), int.as_ref());
+    }
+
     /// validates that the in-memory representation follows [the spec](https://arrow.apache.org/docs/format/Columnar.html#struct-layout)
     #[test]
     fn test_struct_array_from_vec() {

arrow-array/src/record_batch.rs

@@ -20,6 +20,7 @@
 
 use crate::{new_empty_array, Array, ArrayRef, StructArray};
 use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef};
+use std::ops::Index;
 use std::sync::Arc;
 
 /// Trait for types that can read `RecordBatch`'s.
@@ -288,6 +289,13 @@ impl RecordBatch {
         &self.columns[index]
     }
 
+    /// Get a reference to a column's array by name.
+    pub fn column_by_name(&self, name: &str) -> Option<&ArrayRef> {
+        self.schema()
+            .column_with_name(name)
+            .map(|(index, _)| &self.columns[index])
+    }
+
     /// Get a reference to all columns in the record batch.
     pub fn columns(&self) -> &[ArrayRef] {
         &self.columns[..]
@@ -473,6 +481,19 @@ impl From<RecordBatch> for StructArray {
     }
 }
 
+impl Index<&str> for RecordBatch {
+    type Output = ArrayRef;
+
+    /// Get a reference to a column's array by name.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the name is not in the schema.
+    fn index(&self, name: &str) -> &Self::Output {
+        self.column_by_name(name).unwrap()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -746,6 +767,25 @@ mod tests {
         assert_eq!(batch1, batch2);
     }
 
+    /// validates if the record batch can be accessed using `column_name` as index i.e. `record_batch["column_name"]`
+    #[test]
+    fn record_batch_index_access() {
+        let id_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
+        let val_arr = Arc::new(Int32Array::from(vec![5, 6, 7, 8]));
+        let schema1 = Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("val", DataType::Int32, false),
+        ]);
+        let record_batch = RecordBatch::try_new(
+            Arc::new(schema1),
+            vec![id_arr.clone(), val_arr.clone()],
+        )
+        .unwrap();
+
+        assert_eq!(record_batch["id"].as_ref(), id_arr.as_ref());
+        assert_eq!(record_batch["val"].as_ref(), val_arr.as_ref());
+    }
+
     #[test]
     fn record_batch_vals_ne() {
         let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);