1use std::collections::BTreeSet;
2
3use anyhow::{anyhow, Result};
4use polars_arrow::datatypes::ArrowSchema as Schema;
5
6pub fn project_schema(schema: &Schema, field_selection: &BTreeSet<String>) -> Schema {
12 schema
13 .fields
14 .iter()
15 .filter(|f| field_selection.contains(&f.name))
16 .cloned()
17 .collect::<Vec<_>>()
18 .into()
19}
20
21pub fn try_project_schema(schema: &Schema, field_selection: &BTreeSet<String>) -> Result<Schema> {
25 let known: BTreeSet<&String> = schema.fields.iter().map(|f| &f.name).collect();
26 let missing: Vec<&String> = field_selection
27 .iter()
28 .filter(|name| !known.contains(name))
29 .collect();
30 if !missing.is_empty() {
31 return Err(anyhow!(
32 "selected columns not found in schema: {}",
33 missing
34 .into_iter()
35 .map(String::as_str)
36 .collect::<Vec<_>>()
37 .join(", ")
38 ));
39 }
40
41 Ok(project_schema(schema, field_selection))
42}