1use std::collections::BTreeSet;
2
3use polars_arrow::datatypes::ArrowSchema as Schema;
4
5pub fn project_schema(schema: &Schema, field_selection: &BTreeSet<String>) -> Schema {
6 let mut select_indices = Vec::new();
7 for col_name in field_selection.iter() {
8 if let Some((idx, _)) = schema
9 .fields
10 .iter()
11 .enumerate()
12 .find(|(_, f)| &f.name == col_name)
13 {
14 select_indices.push(idx);
15 }
16 }
17
18 let schema: Schema = schema
19 .fields
20 .iter()
21 .filter(|f| field_selection.contains(&f.name))
22 .cloned()
23 .collect::<Vec<_>>()
24 .into();
25
26 schema
27}