iceberg/spec/schema/
visitor.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use super::*;
19use crate::spec::VariantType;
20
21/// A post order schema visitor.
22///
23/// For order of methods called, please refer to [`visit_schema`].
24pub trait SchemaVisitor {
25    /// Return type of this visitor.
26    type T;
27
28    /// Called before struct field.
29    fn before_struct_field(&mut self, _field: &NestedFieldRef) -> Result<()> {
30        Ok(())
31    }
32    /// Called after struct field.
33    fn after_struct_field(&mut self, _field: &NestedFieldRef) -> Result<()> {
34        Ok(())
35    }
36    /// Called before list field.
37    fn before_list_element(&mut self, _field: &NestedFieldRef) -> Result<()> {
38        Ok(())
39    }
40    /// Called after list field.
41    fn after_list_element(&mut self, _field: &NestedFieldRef) -> Result<()> {
42        Ok(())
43    }
44    /// Called before map key field.
45    fn before_map_key(&mut self, _field: &NestedFieldRef) -> Result<()> {
46        Ok(())
47    }
48    /// Called after map key field.
49    fn after_map_key(&mut self, _field: &NestedFieldRef) -> Result<()> {
50        Ok(())
51    }
52    /// Called before map value field.
53    fn before_map_value(&mut self, _field: &NestedFieldRef) -> Result<()> {
54        Ok(())
55    }
56    /// Called after map value field.
57    fn after_map_value(&mut self, _field: &NestedFieldRef) -> Result<()> {
58        Ok(())
59    }
60
61    /// Called after schema's type visited.
62    fn schema(&mut self, schema: &Schema, value: Self::T) -> Result<Self::T>;
63    /// Called after struct's field type visited.
64    fn field(&mut self, field: &NestedFieldRef, value: Self::T) -> Result<Self::T>;
65    /// Called after struct's fields visited.
66    fn r#struct(&mut self, r#struct: &StructType, results: Vec<Self::T>) -> Result<Self::T>;
67    /// Called after list fields visited.
68    fn list(&mut self, list: &ListType, value: Self::T) -> Result<Self::T>;
69    /// Called after map's key and value fields visited.
70    fn map(&mut self, map: &MapType, key_value: Self::T, value: Self::T) -> Result<Self::T>;
71    /// Called when see a primitive type.
72    fn primitive(&mut self, p: &PrimitiveType) -> Result<Self::T>;
73
74    /// Called when see a variant type.
75    fn variant(&mut self, v: &VariantType) -> Result<Self::T>;
76}
77
78/// Visiting a type in post order.
79pub(crate) fn visit_type<V: SchemaVisitor>(r#type: &Type, visitor: &mut V) -> Result<V::T> {
80    match r#type {
81        Type::Primitive(p) => visitor.primitive(p),
82        Type::List(list) => {
83            visitor.before_list_element(&list.element_field)?;
84            let value = visit_type(&list.element_field.field_type, visitor)?;
85            visitor.after_list_element(&list.element_field)?;
86            visitor.list(list, value)
87        }
88        Type::Map(map) => {
89            let key_result = {
90                visitor.before_map_key(&map.key_field)?;
91                let ret = visit_type(&map.key_field.field_type, visitor)?;
92                visitor.after_map_key(&map.key_field)?;
93                ret
94            };
95
96            let value_result = {
97                visitor.before_map_value(&map.value_field)?;
98                let ret = visit_type(&map.value_field.field_type, visitor)?;
99                visitor.after_map_value(&map.value_field)?;
100                ret
101            };
102
103            visitor.map(map, key_result, value_result)
104        }
105        Type::Struct(s) => visit_struct(s, visitor),
106        Type::Variant(v) => visitor.variant(v),
107    }
108}
109
110/// Visit struct type in post order.
111pub fn visit_struct<V: SchemaVisitor>(s: &StructType, visitor: &mut V) -> Result<V::T> {
112    let mut results = Vec::with_capacity(s.fields().len());
113    for field in s.fields() {
114        visitor.before_struct_field(field)?;
115        let result = visit_type(&field.field_type, visitor)?;
116        visitor.after_struct_field(field)?;
117        let result = visitor.field(field, result)?;
118        results.push(result);
119    }
120
121    visitor.r#struct(s, results)
122}
123
124/// Visit schema in post order.
125pub fn visit_schema<V: SchemaVisitor>(schema: &Schema, visitor: &mut V) -> Result<V::T> {
126    let result = visit_struct(&schema.r#struct, visitor)?;
127    visitor.schema(schema, result)
128}
129
130/// A post order schema visitor with partner.
131///
132/// For order of methods called, please refer to [`visit_schema_with_partner`].
133pub trait SchemaWithPartnerVisitor<P> {
134    /// Return type of this visitor.
135    type T;
136
137    /// Called before struct field.
138    fn before_struct_field(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
139        Ok(())
140    }
141    /// Called after struct field.
142    fn after_struct_field(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
143        Ok(())
144    }
145    /// Called before list field.
146    fn before_list_element(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
147        Ok(())
148    }
149    /// Called after list field.
150    fn after_list_element(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
151        Ok(())
152    }
153    /// Called before map key field.
154    fn before_map_key(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
155        Ok(())
156    }
157    /// Called after map key field.
158    fn after_map_key(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
159        Ok(())
160    }
161    /// Called before map value field.
162    fn before_map_value(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
163        Ok(())
164    }
165    /// Called after map value field.
166    fn after_map_value(&mut self, _field: &NestedFieldRef, _partner: &P) -> Result<()> {
167        Ok(())
168    }
169
170    /// Called after schema's type visited.
171    fn schema(&mut self, schema: &Schema, partner: &P, value: Self::T) -> Result<Self::T>;
172    /// Called after struct's field type visited.
173    fn field(&mut self, field: &NestedFieldRef, partner: &P, value: Self::T) -> Result<Self::T>;
174    /// Called after struct's fields visited.
175    fn r#struct(
176        &mut self,
177        r#struct: &StructType,
178        partner: &P,
179        results: Vec<Self::T>,
180    ) -> Result<Self::T>;
181    /// Called after list fields visited.
182    fn list(&mut self, list: &ListType, partner: &P, value: Self::T) -> Result<Self::T>;
183    /// Called after map's key and value fields visited.
184    fn map(
185        &mut self,
186        map: &MapType,
187        partner: &P,
188        key_value: Self::T,
189        value: Self::T,
190    ) -> Result<Self::T>;
191    /// Called when see a primitive type.
192    fn primitive(&mut self, p: &PrimitiveType, partner: &P) -> Result<Self::T>;
193    /// Called when see a variant type.
194    fn variant(&mut self, v: &VariantType, partner: &P) -> Result<Self::T>;
195}
196
197/// Accessor used to get child partner from parent partner.
198pub trait PartnerAccessor<P> {
199    /// Get the struct partner from schema partner.
200    fn struct_partner<'a>(&self, schema_partner: &'a P) -> Result<&'a P>;
201    /// Get the field partner from struct partner.
202    fn field_partner<'a>(&self, struct_partner: &'a P, field: &NestedField) -> Result<&'a P>;
203    /// Get the list element partner from list partner.
204    fn list_element_partner<'a>(&self, list_partner: &'a P) -> Result<&'a P>;
205    /// Get the map key partner from map partner.
206    fn map_key_partner<'a>(&self, map_partner: &'a P) -> Result<&'a P>;
207    /// Get the map value partner from map partner.
208    fn map_value_partner<'a>(&self, map_partner: &'a P) -> Result<&'a P>;
209}
210
211/// Visiting a type in post order.
212pub(crate) fn visit_type_with_partner<P, V: SchemaWithPartnerVisitor<P>, A: PartnerAccessor<P>>(
213    r#type: &Type,
214    partner: &P,
215    visitor: &mut V,
216    accessor: &A,
217) -> Result<V::T> {
218    match r#type {
219        Type::Primitive(p) => visitor.primitive(p, partner),
220        Type::List(list) => {
221            let list_element_partner = accessor.list_element_partner(partner)?;
222            visitor.before_list_element(&list.element_field, list_element_partner)?;
223            let element_results = visit_type_with_partner(
224                &list.element_field.field_type,
225                list_element_partner,
226                visitor,
227                accessor,
228            )?;
229            visitor.after_list_element(&list.element_field, list_element_partner)?;
230            visitor.list(list, partner, element_results)
231        }
232        Type::Map(map) => {
233            let key_partner = accessor.map_key_partner(partner)?;
234            visitor.before_map_key(&map.key_field, key_partner)?;
235            let key_result =
236                visit_type_with_partner(&map.key_field.field_type, key_partner, visitor, accessor)?;
237            visitor.after_map_key(&map.key_field, key_partner)?;
238
239            let value_partner = accessor.map_value_partner(partner)?;
240            visitor.before_map_value(&map.value_field, value_partner)?;
241            let value_result = visit_type_with_partner(
242                &map.value_field.field_type,
243                value_partner,
244                visitor,
245                accessor,
246            )?;
247            visitor.after_map_value(&map.value_field, value_partner)?;
248
249            visitor.map(map, partner, key_result, value_result)
250        }
251        Type::Struct(s) => visit_struct_with_partner(s, partner, visitor, accessor),
252        Type::Variant(v) => visitor.variant(v, partner),
253    }
254}
255
256/// Visit struct type in post order.
257pub fn visit_struct_with_partner<P, V: SchemaWithPartnerVisitor<P>, A: PartnerAccessor<P>>(
258    s: &StructType,
259    partner: &P,
260    visitor: &mut V,
261    accessor: &A,
262) -> Result<V::T> {
263    let mut results = Vec::with_capacity(s.fields().len());
264    for field in s.fields() {
265        let field_partner = accessor.field_partner(partner, field)?;
266        visitor.before_struct_field(field, field_partner)?;
267        let result = visit_type_with_partner(&field.field_type, field_partner, visitor, accessor)?;
268        visitor.after_struct_field(field, field_partner)?;
269        let result = visitor.field(field, field_partner, result)?;
270        results.push(result);
271    }
272
273    visitor.r#struct(s, partner, results)
274}
275
276/// Visit schema in post order.
277pub fn visit_schema_with_partner<P, V: SchemaWithPartnerVisitor<P>, A: PartnerAccessor<P>>(
278    schema: &Schema,
279    partner: &P,
280    visitor: &mut V,
281    accessor: &A,
282) -> Result<V::T> {
283    let result = visit_struct_with_partner(
284        &schema.r#struct,
285        accessor.struct_partner(partner)?,
286        visitor,
287        accessor,
288    )?;
289    visitor.schema(schema, partner, result)
290}
iceberg/spec/schema/visitor.rs

iceberg/spec/schema/
visitor.rs
