vegafusion_core/spec/transform/
stack.rs

1use crate::expression::column_usage::{ColumnUsage, DatasetsColumnUsage, VlSelectionFields};
2use crate::spec::transform::{TransformColumns, TransformSpecTrait};
3use crate::spec::values::{CompareSpec, Field};
4use crate::task_graph::graph::ScopedVariable;
5use crate::task_graph::scope::TaskScope;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::collections::HashMap;
9use vegafusion_common::escape::unescape_field;
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12pub struct StackTransformSpec {
13    pub field: Field,
14
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub groupby: Option<Vec<Field>>,
17
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub sort: Option<CompareSpec>,
20
21    #[serde(rename = "as", skip_serializing_if = "Option::is_none")]
22    pub as_: Option<Vec<String>>,
23
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub offset: Option<StackOffsetSpec>,
26
27    #[serde(flatten)]
28    pub extra: HashMap<String, Value>,
29}
30
31impl StackTransformSpec {
32    pub fn as_(&self) -> Vec<String> {
33        self.as_
34            .clone()
35            .unwrap_or_else(|| vec!["y0".to_string(), "y1".to_string()])
36    }
37
38    pub fn offset(&self) -> StackOffsetSpec {
39        self.offset.clone().unwrap_or(StackOffsetSpec::Zero)
40    }
41}
42
43#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(rename_all = "lowercase")]
45pub enum StackOffsetSpec {
46    Zero,
47    Center,
48    Normalize,
49}
50
51impl TransformSpecTrait for StackTransformSpec {
52    fn transform_columns(
53        &self,
54        datum_var: &Option<ScopedVariable>,
55        _usage_scope: &[u32],
56        _task_scope: &TaskScope,
57        _vl_selection_fields: &VlSelectionFields,
58    ) -> TransformColumns {
59        if let Some(datum_var) = datum_var {
60            // Init column usage with field
61            let mut col_usage = ColumnUsage::from(unescape_field(&self.field.field()).as_str());
62
63            // Add groupby usage
64            if let Some(groupby) = self.groupby.as_ref() {
65                let groupby: Vec<_> = groupby
66                    .iter()
67                    .map(|field| unescape_field(&field.field()))
68                    .collect();
69                col_usage = col_usage.union(&ColumnUsage::from(groupby.as_slice()));
70            }
71
72            // Add sort usage
73            if let Some(compares) = self.sort.as_ref() {
74                let unescaped_sort_fields: Vec<_> = compares
75                    .field
76                    .to_vec()
77                    .iter()
78                    .map(|f| unescape_field(f))
79                    .collect();
80                col_usage = col_usage.union(&ColumnUsage::from(unescaped_sort_fields.as_slice()));
81            }
82
83            // Build produced
84            let produced = ColumnUsage::from(self.as_().as_slice());
85
86            let usage = DatasetsColumnUsage::empty().with_column_usage(datum_var, col_usage);
87            TransformColumns::PassThrough { usage, produced }
88        } else {
89            TransformColumns::Unknown
90        }
91    }
92
93    fn local_datetime_columns_produced(
94        &self,
95        input_local_datetime_columns: &[String],
96    ) -> Vec<String> {
97        // Keep input local datetime columns as stack passes through all input columns and will
98        // never create a local datetime column.
99        Vec::from(input_local_datetime_columns)
100    }
101}
vegafusion_core/spec/transform/stack.rs

vegafusion_core/spec/transform/
stack.rs
