From fa9e6e98b7d88d25ee0e9c6d7fca0125025296d1 Mon Sep 17 00:00:00 2001
From: Alex Wilcoxson <alex.wilcoxson@relativity.com>
Date: Thu, 16 May 2024 13:08:35 -0500
Subject: [PATCH] fix: enable field_with_name to support nested fields with '.'
 delimiter

---
 crates/core/src/kernel/models/schema.rs | 37 ++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/crates/core/src/kernel/models/schema.rs b/crates/core/src/kernel/models/schema.rs
index 675541a812..161de0352a 100644
--- a/crates/core/src/kernel/models/schema.rs
+++ b/crates/core/src/kernel/models/schema.rs
@@ -259,7 +259,19 @@ impl StructType {
 
     /// Returns a reference of a specific [`StructField`] instance selected by name.
     pub fn field_with_name(&self, name: &str) -> Result<&StructField, Error> {
-        Ok(&self.fields[self.index_of(name)?])
+        match name.split_once('.') {
+            Some((parent, children)) => {
+                let parent_field = &self.fields[self.index_of(parent)?];
+                match parent_field.data_type {
+                    DataType::Struct(ref inner) => Ok(inner.field_with_name(children)?),
+                    _ => Err(Error::Schema(format!(
+                        "Field {} is not a struct type",
+                        parent_field.name()
+                    ))),
+                }
+            }
+            None => Ok(&self.fields[self.index_of(name)?]),
+        }
     }
 
     /// Get all invariants in the schemas
@@ -983,4 +995,27 @@ mod tests {
         );
         assert_eq!(get_hash(&field_1), get_hash(&field_2));
     }
+
+    #[test]
+    fn test_field_with_name() {
+        let schema = StructType::new(vec![
+            StructField::new("a", DataType::STRING, true),
+            StructField::new("b", DataType::INTEGER, true),
+        ]);
+        let field = schema.field_with_name("b").unwrap();
+        assert_eq!(*field, StructField::new("b", DataType::INTEGER, true));
+    }
+
+    #[test]
+    fn test_field_with_name_nested() {
+        let nested = StructType::new(vec![StructField::new("a", DataType::BOOLEAN, true)]);
+        let schema = StructType::new(vec![
+            StructField::new("a", DataType::STRING, true),
+            StructField::new("b", DataType::Struct(Box::new(nested)), true),
+        ]);
+
+        let field = schema.field_with_name("b.a").unwrap();
+
+        assert_eq!(*field, StructField::new("a", DataType::BOOLEAN, true));
+    }
 }