vegafusion_core/transform/
mod.rs

1use crate::error::VegaFusionError;
2use crate::proto::gen::tasks::Variable;
3use crate::proto::gen::transforms::transform::TransformKind;
4use crate::proto::gen::transforms::{
5    Aggregate, Bin, Collect, Extent, Filter, Fold, Formula, Identifier, Impute, Pivot, Project,
6    Sequence, Stack, TimeUnit,
7};
8use crate::proto::gen::transforms::{JoinAggregate, Transform, Window};
9use crate::spec::transform::TransformSpec;
10use crate::task_graph::task::InputVariable;
11use std::convert::TryFrom;
12
13pub mod aggregate;
14pub mod bin;
15pub mod collect;
16pub mod extent;
17pub mod filter;
18pub mod fold;
19pub mod formula;
20pub mod identifier;
21pub mod impute;
22pub mod joinaggregate;
23pub mod pipeline;
24pub mod pivot;
25pub mod project;
26pub mod sequence;
27pub mod stack;
28pub mod timeunit;
29pub mod window;
30
31impl TryFrom<&TransformSpec> for TransformKind {
32    type Error = VegaFusionError;
33
34    fn try_from(value: &TransformSpec) -> std::result::Result<Self, Self::Error> {
35        Ok(match value {
36            TransformSpec::Extent(tx_spec) => Self::Extent(Extent::new(tx_spec)),
37            TransformSpec::Filter(tx_spec) => Self::Filter(Filter::try_new(tx_spec)?),
38            TransformSpec::Formula(tx_spec) => Self::Formula(Formula::try_new(tx_spec)?),
39            TransformSpec::Bin(tx_spec) => Self::Bin(Bin::try_new(tx_spec)?),
40            TransformSpec::Aggregate(tx_spec) => Self::Aggregate(Aggregate::new(tx_spec)),
41            TransformSpec::Collect(tx_spec) => Self::Collect(Collect::try_new(tx_spec)?),
42            TransformSpec::Timeunit(tx_spec) => Self::Timeunit(TimeUnit::try_new(tx_spec)?),
43            TransformSpec::JoinAggregate(tx_spec) => {
44                Self::Joinaggregate(JoinAggregate::new(tx_spec))
45            }
46            TransformSpec::Window(tx_spec) => Self::Window(Window::try_new(tx_spec)?),
47            TransformSpec::Project(tx_spec) => Self::Project(Project::try_new(tx_spec)?),
48            TransformSpec::Stack(tx_spec) => Self::Stack(Stack::try_new(tx_spec)?),
49            TransformSpec::Impute(tx_spec) => Self::Impute(Impute::try_new(tx_spec)?),
50            TransformSpec::Pivot(tx_spec) => Self::Pivot(Pivot::try_new(tx_spec)?),
51            TransformSpec::Identifier(tx_spec) => Self::Identifier(Identifier::try_new(tx_spec)?),
52            TransformSpec::Fold(tx_spec) => Self::Fold(Fold::try_new(tx_spec)?),
53            TransformSpec::Sequence(tx_spec) => Self::Sequence(Sequence::try_new(tx_spec)?),
54            _ => {
55                return Err(VegaFusionError::parse(format!(
56                    "Unsupported transform: {value:?}"
57                )))
58            }
59        })
60    }
61}
62
63impl TryFrom<&TransformSpec> for Transform {
64    type Error = VegaFusionError;
65
66    fn try_from(value: &TransformSpec) -> Result<Self, Self::Error> {
67        Ok(Self {
68            transform_kind: Some(TransformKind::try_from(value)?),
69        })
70    }
71}
72
73impl TransformKind {
74    pub fn as_dependencies_trait(&self) -> &dyn TransformDependencies {
75        match self {
76            TransformKind::Filter(tx) => tx,
77            TransformKind::Extent(tx) => tx,
78            TransformKind::Formula(tx) => tx,
79            TransformKind::Bin(tx) => tx,
80            TransformKind::Aggregate(tx) => tx,
81            TransformKind::Collect(tx) => tx,
82            TransformKind::Timeunit(tx) => tx,
83            TransformKind::Joinaggregate(tx) => tx,
84            TransformKind::Window(tx) => tx,
85            TransformKind::Project(tx) => tx,
86            TransformKind::Stack(tx) => tx,
87            TransformKind::Impute(tx) => tx,
88            TransformKind::Pivot(tx) => tx,
89            TransformKind::Identifier(tx) => tx,
90            TransformKind::Fold(tx) => tx,
91            TransformKind::Sequence(tx) => tx,
92        }
93    }
94}
95
96impl Transform {
97    pub fn transform_kind(&self) -> &TransformKind {
98        self.transform_kind.as_ref().unwrap()
99    }
100}
101
102pub trait TransformDependencies: Send + Sync {
103    fn input_vars(&self) -> Vec<InputVariable> {
104        Vec::new()
105    }
106
107    fn output_vars(&self) -> Vec<Variable> {
108        Vec::new()
109    }
110}
111
112impl TransformDependencies for Transform {
113    fn input_vars(&self) -> Vec<InputVariable> {
114        self.transform_kind().as_dependencies_trait().input_vars()
115    }
116
117    fn output_vars(&self) -> Vec<Variable> {
118        self.transform_kind().as_dependencies_trait().output_vars()
119    }
120}