1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

use std::collections::BTreeSet;

use polars_arrow::datatypes::ArrowSchema as Schema;

pub fn project_schema(schema: &Schema, field_selection: &BTreeSet<String>) -> Schema {
    let mut select_indices = Vec::new();
    for col_name in field_selection.iter() {
        if let Some((idx, _)) = schema
            .fields
            .iter()
            .enumerate()
            .find(|(_, f)| &f.name == col_name)
        {
            select_indices.push(idx);
        }
    }

    let schema: Schema = schema
        .fields
        .iter()
        .filter(|f| field_selection.contains(&f.name))
        .cloned()
        .collect::<Vec<_>>()
        .into();

    schema
}