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
28
29
30
31

use std::collections::BTreeSet;

use anyhow::{Context, Result};

use arrow2::datatypes::Schema;

pub fn project_schema(
    schema: &Schema,
    field_selection: &BTreeSet<String>,
) -> Result<Schema, anyhow::Error> {
    let mut select_indices = Vec::new();
    for col_name in field_selection.iter() {
        let (idx, _) = schema
            .fields
            .iter()
            .enumerate()
            .find(|(_, f)| &f.name == col_name)
            .context(format!("couldn't find column {col_name} in schema"))?;
        select_indices.push(idx);
    }

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

    Ok(schema)
}