iceberg/spec/schema/
prune_columns.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
21struct PruneColumn {
22    selected: HashSet<i32>,
23    select_full_types: bool,
24}
25
26/// Visit a schema and returns only the fields selected by id set
27pub fn prune_columns(
28    schema: &Schema,
29    selected: impl IntoIterator<Item = i32>,
30    select_full_types: bool,
31) -> Result<Type> {
32    let mut visitor = PruneColumn::new(HashSet::from_iter(selected), select_full_types);
33    let result = visit_schema(schema, &mut visitor);
34
35    match result {
36        Ok(s) => {
37            if let Some(struct_type) = s {
38                Ok(struct_type)
39            } else {
40                Ok(Type::Struct(StructType::default()))
41            }
42        }
43        Err(e) => Err(e),
44    }
45}
46
47impl PruneColumn {
48    fn new(selected: HashSet<i32>, select_full_types: bool) -> Self {
49        Self {
50            selected,
51            select_full_types,
52        }
53    }
54
55    fn project_selected_struct(projected_field: Option<Type>) -> Result<StructType> {
56        match projected_field {
57            // If the field is a StructType, return it as such
58            Some(Type::Struct(s)) => Ok(s),
59            Some(_) => Err(Error::new(
60                ErrorKind::Unexpected,
61                "Projected field with struct type must be struct".to_string(),
62            )),
63            // If projected_field is None or not a StructType, return an empty StructType
64            None => Ok(StructType::default()),
65        }
66    }
67    fn project_list(list: &ListType, element_result: Type) -> Result<ListType> {
68        if *list.element_field.field_type == element_result {
69            return Ok(list.clone());
70        }
71        Ok(ListType {
72            element_field: Arc::new(NestedField {
73                id: list.element_field.id,
74                name: list.element_field.name.clone(),
75                required: list.element_field.required,
76                field_type: Box::new(element_result),
77                doc: list.element_field.doc.clone(),
78                initial_default: list.element_field.initial_default.clone(),
79                write_default: list.element_field.write_default.clone(),
80            }),
81        })
82    }
83    fn project_map(map: &MapType, value_result: Type) -> Result<MapType> {
84        if *map.value_field.field_type == value_result {
85            return Ok(map.clone());
86        }
87        Ok(MapType {
88            key_field: map.key_field.clone(),
89            value_field: Arc::new(NestedField {
90                id: map.value_field.id,
91                name: map.value_field.name.clone(),
92                required: map.value_field.required,
93                field_type: Box::new(value_result),
94                doc: map.value_field.doc.clone(),
95                initial_default: map.value_field.initial_default.clone(),
96                write_default: map.value_field.write_default.clone(),
97            }),
98        })
99    }
100}
101
102impl SchemaVisitor for PruneColumn {
103    type T = Option<Type>;
104
105    fn schema(&mut self, _schema: &Schema, value: Option<Type>) -> Result<Option<Type>> {
106        Ok(Some(value.unwrap()))
107    }
108
109    fn field(&mut self, field: &NestedFieldRef, value: Option<Type>) -> Result<Option<Type>> {
110        if self.selected.contains(&field.id) {
111            if self.select_full_types {
112                Ok(Some(*field.field_type.clone()))
113            } else if field.field_type.is_struct() {
114                Ok(Some(Type::Struct(PruneColumn::project_selected_struct(
115                    value,
116                )?)))
117            } else if !field.field_type.is_nested() {
118                Ok(Some(*field.field_type.clone()))
119            } else {
120                Err(Error::new(
121                    ErrorKind::DataInvalid,
122                    "Can't project list or map field directly when not selecting full type."
123                        .to_string(),
124                )
125                .with_context("field_id", field.id.to_string())
126                .with_context("field_type", field.field_type.to_string()))
127            }
128        } else {
129            Ok(value)
130        }
131    }
132
133    fn r#struct(
134        &mut self,
135        r#struct: &StructType,
136        results: Vec<Option<Type>>,
137    ) -> Result<Option<Type>> {
138        let fields = r#struct.fields();
139        let mut selected_field = Vec::with_capacity(fields.len());
140        let mut same_type = true;
141
142        for (field, projected_type) in zip_eq(fields.iter(), results.iter()) {
143            if let Some(projected_type) = projected_type {
144                if *field.field_type == *projected_type {
145                    selected_field.push(field.clone());
146                } else {
147                    same_type = false;
148                    let new_field = NestedField {
149                        id: field.id,
150                        name: field.name.clone(),
151                        required: field.required,
152                        field_type: Box::new(projected_type.clone()),
153                        doc: field.doc.clone(),
154                        initial_default: field.initial_default.clone(),
155                        write_default: field.write_default.clone(),
156                    };
157                    selected_field.push(Arc::new(new_field));
158                }
159            }
160        }
161
162        if !selected_field.is_empty() {
163            if selected_field.len() == fields.len() && same_type {
164                return Ok(Some(Type::Struct(r#struct.clone())));
165            } else {
166                return Ok(Some(Type::Struct(StructType::new(selected_field))));
167            }
168        }
169        Ok(None)
170    }
171
172    fn list(&mut self, list: &ListType, value: Option<Type>) -> Result<Option<Type>> {
173        if self.selected.contains(&list.element_field.id) {
174            if self.select_full_types {
175                Ok(Some(Type::List(list.clone())))
176            } else if list.element_field.field_type.is_struct() {
177                let projected_struct = PruneColumn::project_selected_struct(value).unwrap();
178                Ok(Some(Type::List(PruneColumn::project_list(
179                    list,
180                    Type::Struct(projected_struct),
181                )?)))
182            } else if list.element_field.field_type.is_primitive() {
183                Ok(Some(Type::List(list.clone())))
184            } else {
185                Err(Error::new(
186                    ErrorKind::DataInvalid,
187                    format!(
188                        "Cannot explicitly project List or Map types, List element {} of type {} was selected",
189                        list.element_field.id, list.element_field.field_type
190                    ),
191                ))
192            }
193        } else if let Some(result) = value {
194            Ok(Some(Type::List(PruneColumn::project_list(list, result)?)))
195        } else {
196            Ok(None)
197        }
198    }
199
200    fn map(
201        &mut self,
202        map: &MapType,
203        _key_value: Option<Type>,
204        value: Option<Type>,
205    ) -> Result<Option<Type>> {
206        if self.selected.contains(&map.value_field.id) {
207            if self.select_full_types {
208                Ok(Some(Type::Map(map.clone())))
209            } else if map.value_field.field_type.is_struct() {
210                let projected_struct =
211                    PruneColumn::project_selected_struct(Some(value.unwrap())).unwrap();
212                Ok(Some(Type::Map(PruneColumn::project_map(
213                    map,
214                    Type::Struct(projected_struct),
215                )?)))
216            } else if map.value_field.field_type.is_primitive() {
217                Ok(Some(Type::Map(map.clone())))
218            } else {
219                Err(Error::new(
220                    ErrorKind::DataInvalid,
221                    format!(
222                        "Cannot explicitly project List or Map types, Map value {} of type {} was selected",
223                        map.value_field.id, map.value_field.field_type
224                    ),
225                ))
226            }
227        } else if let Some(value_result) = value {
228            Ok(Some(Type::Map(PruneColumn::project_map(
229                map,
230                value_result,
231            )?)))
232        } else if self.selected.contains(&map.key_field.id) {
233            Ok(Some(Type::Map(map.clone())))
234        } else {
235            Ok(None)
236        }
237    }
238
239    fn primitive(&mut self, _p: &PrimitiveType) -> Result<Option<Type>> {
240        Ok(None)
241    }
242
243    fn variant(&mut self, _v: &VariantType) -> Result<Self::T> {
244        Ok(None)
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use Type::Primitive;
251
252    use super::*;
253    use crate::spec::schema::tests::table_schema_nested;
254
255    #[test]
256    fn test_schema_prune_columns_string() {
257        let expected_type = Type::from(
258            Schema::builder()
259                .with_fields(vec![
260                    NestedField::optional(1, "foo", Primitive(PrimitiveType::String)).into(),
261                ])
262                .build()
263                .unwrap()
264                .as_struct()
265                .clone(),
266        );
267        let schema = table_schema_nested();
268        let selected: HashSet<i32> = HashSet::from([1]);
269        let result = prune_columns(&schema, selected, false);
270        assert!(result.is_ok());
271        assert_eq!(result.unwrap(), expected_type);
272    }
273
274    #[test]
275    fn test_schema_prune_columns_string_full() {
276        let expected_type = Type::from(
277            Schema::builder()
278                .with_fields(vec![
279                    NestedField::optional(1, "foo", Primitive(PrimitiveType::String)).into(),
280                ])
281                .build()
282                .unwrap()
283                .as_struct()
284                .clone(),
285        );
286        let schema = table_schema_nested();
287        let selected: HashSet<i32> = HashSet::from([1]);
288        let result = prune_columns(&schema, selected, true);
289        assert!(result.is_ok());
290        assert_eq!(result.unwrap(), expected_type);
291    }
292
293    #[test]
294    fn test_schema_prune_columns_list() {
295        let expected_type = Type::from(
296            Schema::builder()
297                .with_fields(vec![
298                    NestedField::required(
299                        4,
300                        "qux",
301                        Type::List(ListType {
302                            element_field: NestedField::list_element(
303                                5,
304                                Primitive(PrimitiveType::String),
305                                true,
306                            )
307                            .into(),
308                        }),
309                    )
310                    .into(),
311                ])
312                .build()
313                .unwrap()
314                .as_struct()
315                .clone(),
316        );
317        let schema = table_schema_nested();
318        let selected: HashSet<i32> = HashSet::from([5]);
319        let result = prune_columns(&schema, selected, false);
320        assert!(result.is_ok());
321        assert_eq!(result.unwrap(), expected_type);
322    }
323
324    #[test]
325    fn test_prune_columns_list_itself() {
326        let schema = table_schema_nested();
327        let selected: HashSet<i32> = HashSet::from([4]);
328        let result = prune_columns(&schema, selected, false);
329        assert!(result.is_err());
330    }
331
332    #[test]
333    fn test_schema_prune_columns_list_full() {
334        let expected_type = Type::from(
335            Schema::builder()
336                .with_fields(vec![
337                    NestedField::required(
338                        4,
339                        "qux",
340                        Type::List(ListType {
341                            element_field: NestedField::list_element(
342                                5,
343                                Primitive(PrimitiveType::String),
344                                true,
345                            )
346                            .into(),
347                        }),
348                    )
349                    .into(),
350                ])
351                .build()
352                .unwrap()
353                .as_struct()
354                .clone(),
355        );
356        let schema = table_schema_nested();
357        let selected: HashSet<i32> = HashSet::from([5]);
358        let result = prune_columns(&schema, selected, true);
359        assert!(result.is_ok());
360        assert_eq!(result.unwrap(), expected_type);
361    }
362
363    #[test]
364    fn test_prune_columns_map() {
365        let expected_type = Type::from(
366            Schema::builder()
367                .with_fields(vec![
368                    NestedField::required(
369                        6,
370                        "quux",
371                        Type::Map(MapType {
372                            key_field: NestedField::map_key_element(
373                                7,
374                                Primitive(PrimitiveType::String),
375                            )
376                            .into(),
377                            value_field: NestedField::map_value_element(
378                                8,
379                                Type::Map(MapType {
380                                    key_field: NestedField::map_key_element(
381                                        9,
382                                        Primitive(PrimitiveType::String),
383                                    )
384                                    .into(),
385                                    value_field: NestedField::map_value_element(
386                                        10,
387                                        Primitive(PrimitiveType::Int),
388                                        true,
389                                    )
390                                    .into(),
391                                }),
392                                true,
393                            )
394                            .into(),
395                        }),
396                    )
397                    .into(),
398                ])
399                .build()
400                .unwrap()
401                .as_struct()
402                .clone(),
403        );
404        let schema = table_schema_nested();
405        let selected: HashSet<i32> = HashSet::from([9]);
406        let result = prune_columns(&schema, selected, false);
407        assert!(result.is_ok());
408        assert_eq!(result.unwrap(), expected_type);
409    }
410
411    #[test]
412    fn test_prune_columns_map_itself() {
413        let schema = table_schema_nested();
414        let selected: HashSet<i32> = HashSet::from([6]);
415        let result = prune_columns(&schema, selected, false);
416        assert!(result.is_err());
417    }
418
419    #[test]
420    fn test_prune_columns_map_full() {
421        let expected_type = Type::from(
422            Schema::builder()
423                .with_fields(vec![
424                    NestedField::required(
425                        6,
426                        "quux",
427                        Type::Map(MapType {
428                            key_field: NestedField::map_key_element(
429                                7,
430                                Primitive(PrimitiveType::String),
431                            )
432                            .into(),
433                            value_field: NestedField::map_value_element(
434                                8,
435                                Type::Map(MapType {
436                                    key_field: NestedField::map_key_element(
437                                        9,
438                                        Primitive(PrimitiveType::String),
439                                    )
440                                    .into(),
441                                    value_field: NestedField::map_value_element(
442                                        10,
443                                        Primitive(PrimitiveType::Int),
444                                        true,
445                                    )
446                                    .into(),
447                                }),
448                                true,
449                            )
450                            .into(),
451                        }),
452                    )
453                    .into(),
454                ])
455                .build()
456                .unwrap()
457                .as_struct()
458                .clone(),
459        );
460        let schema = table_schema_nested();
461        let selected: HashSet<i32> = HashSet::from([9]);
462        let result = prune_columns(&schema, selected, true);
463        assert!(result.is_ok());
464        assert_eq!(result.unwrap(), expected_type);
465    }
466
467    #[test]
468    fn test_prune_columns_map_key() {
469        let expected_type = Type::from(
470            Schema::builder()
471                .with_fields(vec![
472                    NestedField::required(
473                        6,
474                        "quux",
475                        Type::Map(MapType {
476                            key_field: NestedField::map_key_element(
477                                7,
478                                Primitive(PrimitiveType::String),
479                            )
480                            .into(),
481                            value_field: NestedField::map_value_element(
482                                8,
483                                Type::Map(MapType {
484                                    key_field: NestedField::map_key_element(
485                                        9,
486                                        Primitive(PrimitiveType::String),
487                                    )
488                                    .into(),
489                                    value_field: NestedField::map_value_element(
490                                        10,
491                                        Primitive(PrimitiveType::Int),
492                                        true,
493                                    )
494                                    .into(),
495                                }),
496                                true,
497                            )
498                            .into(),
499                        }),
500                    )
501                    .into(),
502                ])
503                .build()
504                .unwrap()
505                .as_struct()
506                .clone(),
507        );
508        let schema = table_schema_nested();
509        let selected: HashSet<i32> = HashSet::from([10]);
510        let result = prune_columns(&schema, selected, false);
511        assert!(result.is_ok());
512        assert_eq!(result.unwrap(), expected_type);
513    }
514
515    #[test]
516    fn test_prune_columns_struct() {
517        let expected_type = Type::from(
518            Schema::builder()
519                .with_fields(vec![
520                    NestedField::optional(
521                        15,
522                        "person",
523                        Type::Struct(StructType::new(vec![
524                            NestedField::optional(16, "name", Primitive(PrimitiveType::String))
525                                .into(),
526                        ])),
527                    )
528                    .into(),
529                ])
530                .build()
531                .unwrap()
532                .as_struct()
533                .clone(),
534        );
535        let schema = table_schema_nested();
536        let selected: HashSet<i32> = HashSet::from([16]);
537        let result = prune_columns(&schema, selected, false);
538        assert!(result.is_ok());
539        assert_eq!(result.unwrap(), expected_type);
540    }
541
542    #[test]
543    fn test_prune_columns_struct_full() {
544        let expected_type = Type::from(
545            Schema::builder()
546                .with_fields(vec![
547                    NestedField::optional(
548                        15,
549                        "person",
550                        Type::Struct(StructType::new(vec![
551                            NestedField::optional(16, "name", Primitive(PrimitiveType::String))
552                                .into(),
553                        ])),
554                    )
555                    .into(),
556                ])
557                .build()
558                .unwrap()
559                .as_struct()
560                .clone(),
561        );
562        let schema = table_schema_nested();
563        let selected: HashSet<i32> = HashSet::from([16]);
564        let result = prune_columns(&schema, selected, true);
565        assert!(result.is_ok());
566        assert_eq!(result.unwrap(), expected_type);
567    }
568
569    #[test]
570    fn test_prune_columns_empty_struct() {
571        let schema_with_empty_struct_field = Schema::builder()
572            .with_fields(vec![
573                NestedField::optional(15, "person", Type::Struct(StructType::new(vec![]))).into(),
574            ])
575            .build()
576            .unwrap();
577        let expected_type = Type::from(
578            Schema::builder()
579                .with_fields(vec![
580                    NestedField::optional(15, "person", Type::Struct(StructType::new(vec![])))
581                        .into(),
582                ])
583                .build()
584                .unwrap()
585                .as_struct()
586                .clone(),
587        );
588        let selected: HashSet<i32> = HashSet::from([15]);
589        let result = prune_columns(&schema_with_empty_struct_field, selected, false);
590        assert!(result.is_ok());
591        assert_eq!(result.unwrap(), expected_type);
592    }
593
594    #[test]
595    fn test_prune_columns_empty_struct_full() {
596        let schema_with_empty_struct_field = Schema::builder()
597            .with_fields(vec![
598                NestedField::optional(15, "person", Type::Struct(StructType::new(vec![]))).into(),
599            ])
600            .build()
601            .unwrap();
602        let expected_type = Type::from(
603            Schema::builder()
604                .with_fields(vec![
605                    NestedField::optional(15, "person", Type::Struct(StructType::new(vec![])))
606                        .into(),
607                ])
608                .build()
609                .unwrap()
610                .as_struct()
611                .clone(),
612        );
613        let selected: HashSet<i32> = HashSet::from([15]);
614        let result = prune_columns(&schema_with_empty_struct_field, selected, true);
615        assert!(result.is_ok());
616        assert_eq!(result.unwrap(), expected_type);
617    }
618
619    #[test]
620    fn test_prune_columns_struct_in_map() {
621        let schema_with_struct_in_map_field = Schema::builder()
622            .with_schema_id(1)
623            .with_fields(vec![
624                NestedField::required(
625                    6,
626                    "id_to_person",
627                    Type::Map(MapType {
628                        key_field: NestedField::map_key_element(7, Primitive(PrimitiveType::Int))
629                            .into(),
630                        value_field: NestedField::map_value_element(
631                            8,
632                            Type::Struct(StructType::new(vec![
633                                NestedField::optional(10, "name", Primitive(PrimitiveType::String))
634                                    .into(),
635                                NestedField::required(11, "age", Primitive(PrimitiveType::Int))
636                                    .into(),
637                            ])),
638                            true,
639                        )
640                        .into(),
641                    }),
642                )
643                .into(),
644            ])
645            .build()
646            .unwrap();
647        let expected_type = Type::from(
648            Schema::builder()
649                .with_fields(vec![
650                    NestedField::required(
651                        6,
652                        "id_to_person",
653                        Type::Map(MapType {
654                            key_field: NestedField::map_key_element(
655                                7,
656                                Primitive(PrimitiveType::Int),
657                            )
658                            .into(),
659                            value_field: NestedField::map_value_element(
660                                8,
661                                Type::Struct(StructType::new(vec![
662                                    NestedField::required(11, "age", Primitive(PrimitiveType::Int))
663                                        .into(),
664                                ])),
665                                true,
666                            )
667                            .into(),
668                        }),
669                    )
670                    .into(),
671                ])
672                .build()
673                .unwrap()
674                .as_struct()
675                .clone(),
676        );
677        let selected: HashSet<i32> = HashSet::from([11]);
678        let result = prune_columns(&schema_with_struct_in_map_field, selected, false);
679        assert!(result.is_ok());
680        assert_eq!(result.unwrap(), expected_type);
681    }
682    #[test]
683    fn test_prune_columns_struct_in_map_full() {
684        let schema = Schema::builder()
685            .with_schema_id(1)
686            .with_fields(vec![
687                NestedField::required(
688                    6,
689                    "id_to_person",
690                    Type::Map(MapType {
691                        key_field: NestedField::map_key_element(7, Primitive(PrimitiveType::Int))
692                            .into(),
693                        value_field: NestedField::map_value_element(
694                            8,
695                            Type::Struct(StructType::new(vec![
696                                NestedField::optional(10, "name", Primitive(PrimitiveType::String))
697                                    .into(),
698                                NestedField::required(11, "age", Primitive(PrimitiveType::Int))
699                                    .into(),
700                            ])),
701                            true,
702                        )
703                        .into(),
704                    }),
705                )
706                .into(),
707            ])
708            .build()
709            .unwrap();
710        let expected_type = Type::from(
711            Schema::builder()
712                .with_fields(vec![
713                    NestedField::required(
714                        6,
715                        "id_to_person",
716                        Type::Map(MapType {
717                            key_field: NestedField::map_key_element(
718                                7,
719                                Primitive(PrimitiveType::Int),
720                            )
721                            .into(),
722                            value_field: NestedField::map_value_element(
723                                8,
724                                Type::Struct(StructType::new(vec![
725                                    NestedField::required(11, "age", Primitive(PrimitiveType::Int))
726                                        .into(),
727                                ])),
728                                true,
729                            )
730                            .into(),
731                        }),
732                    )
733                    .into(),
734                ])
735                .build()
736                .unwrap()
737                .as_struct()
738                .clone(),
739        );
740        let selected: HashSet<i32> = HashSet::from([11]);
741        let result = prune_columns(&schema, selected, true);
742        assert!(result.is_ok());
743        assert_eq!(result.unwrap(), expected_type);
744    }
745
746    #[test]
747    fn test_prune_columns_select_original_schema() {
748        let schema = table_schema_nested();
749        let selected: HashSet<i32> = (0..schema.highest_field_id() + 1).collect();
750        let result = prune_columns(&schema, selected, true);
751        assert!(result.is_ok());
752        assert_eq!(result.unwrap(), Type::Struct(schema.as_struct().clone()));
753    }
754
755    #[test]
756    fn test_prune_columns_variant() {
757        // foo (String, id=1) + v (Variant, id=2).
758        let schema = Schema::builder()
759            .with_fields(vec![
760                NestedField::optional(1, "foo", Primitive(PrimitiveType::String)).into(),
761                NestedField::optional(2, "v", Type::Variant(VariantType)).into(),
762            ])
763            .build()
764            .unwrap();
765
766        // A variant is a leaf (like a primitive): selecting it keeps it, the same way
767        // for select_full_types true and false.
768        let only_variant = Type::Struct(StructType::new(vec![
769            NestedField::optional(2, "v", Type::Variant(VariantType)).into(),
770        ]));
771        for full in [false, true] {
772            let result = prune_columns(&schema, HashSet::from([2]), full).unwrap();
773            assert_eq!(result, only_variant, "select_full_types={full}");
774        }
775
776        // Selecting a sibling prunes the variant out.
777        let only_foo = Type::Struct(StructType::new(vec![
778            NestedField::optional(1, "foo", Primitive(PrimitiveType::String)).into(),
779        ]));
780        let result = prune_columns(&schema, HashSet::from([1]), false).unwrap();
781        assert_eq!(result, only_foo);
782    }
783}
iceberg/spec/schema/prune_columns.rs

iceberg/spec/schema/
prune_columns.rs
