vegafusion_core/spec/transform/
project.rs

1use crate::spec::transform::{TransformColumns, TransformSpecTrait};
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use std::collections::HashMap;
5use vegafusion_common::escape::unescape_field;
6
7use crate::error::Result;
8use crate::expression::column_usage::{ColumnUsage, DatasetsColumnUsage, VlSelectionFields};
9use crate::task_graph::graph::ScopedVariable;
10use crate::task_graph::scope::TaskScope;
11use crate::task_graph::task::InputVariable;
12
13/// Struct that serializes to Vega spec for the filter transform
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
15pub struct ProjectTransformSpec {
16    pub fields: Vec<String>,
17
18    #[serde(flatten)]
19    pub extra: HashMap<String, Value>,
20}
21
22impl TransformSpecTrait for ProjectTransformSpec {
23    fn supported(&self) -> bool {
24        true
25    }
26
27    fn input_vars(&self) -> Result<Vec<InputVariable>> {
28        Ok(Default::default())
29    }
30
31    fn transform_columns(
32        &self,
33        datum_var: &Option<ScopedVariable>,
34        _usage_scope: &[u32],
35        _task_scope: &TaskScope,
36        _vl_selection_fields: &VlSelectionFields,
37    ) -> TransformColumns {
38        if let Some(datum_var) = datum_var {
39            let unescaped_fields: Vec<_> = self.fields.iter().map(|f| unescape_field(f)).collect();
40            let col_usage = ColumnUsage::from(unescaped_fields.as_slice());
41            let usage =
42                DatasetsColumnUsage::empty().with_column_usage(datum_var, col_usage.clone());
43            TransformColumns::Overwrite {
44                usage,
45                produced: col_usage,
46            }
47        } else {
48            TransformColumns::Unknown
49        }
50    }
51
52    fn local_datetime_columns_produced(
53        &self,
54        input_local_datetime_columns: &[String],
55    ) -> Vec<String> {
56        // Keep input local datetime columns that are used as projection fields
57        self.fields
58            .iter()
59            .filter_map(|project_field| {
60                if input_local_datetime_columns.contains(&unescape_field(project_field)) {
61                    Some(project_field.clone())
62                } else {
63                    None
64                }
65            })
66            .collect::<Vec<_>>()
67    }
68}