diff --git a/crates/iceberg/src/spec/schema/mod.rs b/crates/iceberg/src/spec/schema/mod.rs index 092fa2516..b95244f42 100644 --- a/crates/iceberg/src/spec/schema/mod.rs +++ b/crates/iceberg/src/spec/schema/mod.rs @@ -23,7 +23,7 @@ use std::sync::Arc; mod utils; mod visitor; -pub use self::visitor::{visit_schema, visit_struct, visit_type, SchemaVisitor}; +pub use self::visitor::*; pub(super) mod _serde; mod id_reassigner; mod index; diff --git a/crates/iceberg/src/spec/schema/visitor.rs b/crates/iceberg/src/spec/schema/visitor.rs index 4f9cf4ebf..0008635bf 100644 --- a/crates/iceberg/src/spec/schema/visitor.rs +++ b/crates/iceberg/src/spec/schema/visitor.rs @@ -162,12 +162,6 @@ pub trait SchemaWithPartnerVisitor

{ Ok(()) } - /// Called before every type, if this function return `Some`, the following visiting will be skipped. - /// This function used to implement early return. - fn visit_type_before(&mut self, _ty: &Type, _partner: &P) -> Result> { - return Ok(None); - } - /// Called after schema's type visited. fn schema(&mut self, schema: &Schema, partner: &P, value: Self::T) -> Result; /// Called after struct's field type visited. @@ -180,14 +174,14 @@ pub trait SchemaWithPartnerVisitor

{ results: Vec, ) -> Result; /// Called after list fields visited. - fn list(&mut self, list: &ListType, partner: &P, value: Vec) -> Result; + fn list(&mut self, list: &ListType, partner: &P, value: Self::T) -> Result; /// Called after map's key and value fields visited. fn map( &mut self, map: &MapType, partner: &P, - key_value: Vec, - value: Vec, + key_value: Self::T, + value: Self::T, ) -> Result; /// Called when see a primitive type. fn primitive(&mut self, p: &PrimitiveType, partner: &P) -> Result; @@ -195,32 +189,16 @@ pub trait SchemaWithPartnerVisitor

{ /// Accessor used to get child partner from parent partner. pub trait PartnerAccessor

{ - /// List partner iterator. - type L: ListPartnerIterator

; - /// Map partner iterator. - type M: MapPartnerIterator

; - /// Get the struct partner from schema partner. fn struct_parner<'a>(&self, schema_partner: &'a P) -> Result<&'a P>; /// Get the field partner from struct partner. - fn field_partner<'a>(&self, struct_partner: &'a P, field_id: i32, field: &str) - -> Result<&'a P>; + fn field_partner<'a>(&self, struct_partner: &'a P, field: &NestedField) -> Result<&'a P>; /// Get the list element partner from list partner. - fn list_element_partner<'a>(&self, list_partner: &'a P) -> Result; + fn list_element_partner<'a>(&self, list_partner: &'a P) -> Result<&'a P>; /// Get the map key partner from map partner. - fn map_element_partner<'a>(&self, map_partner: &'a P) -> Result; -} - -/// Iterator for list partner. -pub trait ListPartnerIterator

{ - /// Get the next partner. - fn next(&mut self) -> Option

; -} - -/// Iterator for map partner. -pub trait MapPartnerIterator

{ - /// Get the next partner. - fn next(&mut self) -> Option<(P, P)>; + fn map_key_partner<'a>(&self, map_partner: &'a P) -> Result<&'a P>; + /// Get the map value partner from map partner. + fn map_value_partner<'a>(&self, map_partner: &'a P) -> Result<&'a P>; } /// Visiting a type in post order. @@ -230,61 +208,38 @@ pub fn visit_type_with_partner, A: PartnerAcce visitor: &mut V, accessor: &A, ) -> Result { - if let Some(res) = visitor.visit_type_before(r#type, partner)? { - return Ok(res); - } match r#type { Type::Primitive(p) => visitor.primitive(p, partner), Type::List(list) => { - let mut results = Vec::new(); - let mut list_element_partner_iter = accessor.list_element_partner(partner)?; - if let Some(list_element_partner) = list_element_partner_iter.next() { - visitor.before_list_element(&list.element_field, &list_element_partner)?; - let value = visit_type_with_partner( - &list.element_field.field_type, - &list_element_partner, - visitor, - accessor, - )?; - visitor.after_list_element(&list.element_field, &list_element_partner)?; - results.push(value); - } - visitor.list(list, partner, results) + let list_element_partner = accessor.list_element_partner(partner)?; + visitor.before_list_element(&list.element_field, list_element_partner)?; + let element_results = visit_type_with_partner( + &list.element_field.field_type, + list_element_partner, + visitor, + accessor, + )?; + visitor.after_list_element(&list.element_field, list_element_partner)?; + visitor.list(list, partner, element_results) } Type::Map(map) => { - let mut k_results = Vec::new(); - let mut v_results = Vec::new(); - let mut kv_partner_iter = accessor.map_element_partner(partner)?; - if let Some((k_partner, v_partner)) = kv_partner_iter.next() { - let key_result = { - visitor.before_map_key(&map.key_field, &k_partner)?; - let ret = visit_type_with_partner( - &map.key_field.field_type, - &k_partner, - visitor, - accessor, - )?; - visitor.after_map_key(&map.key_field, &k_partner)?; - ret - }; - - let value_result = { - visitor.before_map_value(&map.value_field, &v_partner)?; - let ret = visit_type_with_partner( - &map.value_field.field_type, - &v_partner, - visitor, - accessor, - )?; - visitor.after_map_value(&map.value_field, &v_partner)?; - ret - }; + let key_partner = accessor.map_key_partner(partner)?; + visitor.before_map_key(&map.key_field, key_partner)?; + let key_result = + visit_type_with_partner(&map.key_field.field_type, key_partner, visitor, accessor)?; + visitor.after_map_key(&map.key_field, key_partner)?; - k_results.push(key_result); - v_results.push(value_result); - } + let value_partner = accessor.map_value_partner(partner)?; + visitor.before_map_value(&map.value_field, value_partner)?; + let value_result = visit_type_with_partner( + &map.value_field.field_type, + value_partner, + visitor, + accessor, + )?; + visitor.after_map_value(&map.value_field, value_partner)?; - visitor.map(map, partner, k_results, v_results) + visitor.map(map, partner, key_result, value_result) } Type::Struct(s) => visit_struct_with_partner(s, partner, visitor, accessor), } @@ -297,12 +252,9 @@ pub fn visit_struct_with_partner, A: PartnerAc visitor: &mut V, accessor: &A, ) -> Result { - if let Some(res) = visitor.visit_type_before(&Type::Struct(s.clone()), partner)? { - return Ok(res); - } let mut results = Vec::with_capacity(s.fields().len()); for field in s.fields() { - let field_partner = accessor.field_partner(partner, field.id, &field.name)?; + let field_partner = accessor.field_partner(partner, field)?; visitor.before_struct_field(field, field_partner)?; let result = visit_type_with_partner(&field.field_type, field_partner, visitor, accessor)?; visitor.after_struct_field(field, field_partner)?; @@ -327,4 +279,4 @@ pub fn visit_schema_with_partner, A: PartnerAc accessor, )?; visitor.schema(schema, partner, result) -} \ No newline at end of file +}