vegafusion_core/spec/transform/
pivot.rs

1use crate::expression::column_usage::{ColumnUsage, DatasetsColumnUsage, VlSelectionFields};
2use crate::spec::transform::aggregate::AggregateOpSpec;
3use crate::spec::transform::{TransformColumns, TransformSpecTrait};
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/// Struct that serializes to Vega spec for the filter transform
12#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct PivotTransformSpec {
14    pub field: String,
15    pub value: String,
16
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub groupby: Option<Vec<String>>,
19
20    #[serde(skip_serializing_if = "Option::is_none")]
21    pub limit: Option<i32>,
22
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub op: Option<AggregateOpSpec>,
25
26    #[serde(flatten)]
27    pub extra: HashMap<String, Value>,
28}
29
30impl TransformSpecTrait for PivotTransformSpec {
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            // Unknown produced columns
40            let produced = ColumnUsage::Unknown;
41
42            // Compute used columns (groupby, value, and pivot)
43            let mut usage_cols: Vec<_> = self
44                .groupby
45                .clone()
46                .unwrap_or_default()
47                .iter()
48                .map(|f| unescape_field(f))
49                .collect();
50            usage_cols.push(unescape_field(&self.field.clone()));
51            usage_cols.push(unescape_field(&self.value.clone()));
52            let col_usage = ColumnUsage::from(usage_cols.as_slice());
53            let usage = DatasetsColumnUsage::empty().with_column_usage(datum_var, col_usage);
54            TransformColumns::Overwrite { usage, produced }
55        } else {
56            TransformColumns::Unknown
57        }
58    }
59
60    fn local_datetime_columns_produced(
61        &self,
62        input_local_datetime_columns: &[String],
63    ) -> Vec<String> {
64        // Keep input local datetime columns that are used as grouping fields
65        self.groupby
66            .clone()
67            .unwrap_or_default()
68            .iter()
69            .filter_map(|groupby_field| {
70                let unescaped = unescape_field(groupby_field);
71                if input_local_datetime_columns.contains(&unescaped) {
72                    Some(unescaped)
73                } else {
74                    None
75                }
76            })
77            .collect::<Vec<_>>()
78    }
79}
vegafusion_core/spec/transform/pivot.rs

vegafusion_core/spec/transform/
pivot.rs
