Skip to main content

hyperfuel_schema/
util.rs

1use std::collections::BTreeSet;
2
3use anyhow::{anyhow, Result};
4use polars_arrow::datatypes::ArrowSchema as Schema;
5
6/// Project a schema down to the named fields. Missing field names are silently
7/// dropped from the result.
8///
9/// Prefer [`try_project_schema`], which returns an error listing any unknown
10/// names. This non-erroring variant is retained for backward compatibility.
11pub 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
21/// Project a schema down to the named fields, returning an error if any
22/// requested field is not present in the source schema. Use this to catch
23/// typos in field selections instead of silently receiving a partial schema.
24pub 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}