vegafusion_runtime/expression/compiler/builtin_functions/date_time/
date_parts.rs

1use crate::expression::compiler::call::TzTransformFn;
2use crate::expression::compiler::utils::ExprHelpers;
3use crate::task_graph::timezone::RuntimeTzConfig;
4use crate::transform::timeunit::to_timestamp_col;
5use datafusion_expr::{lit, Expr};
6use datafusion_functions::expr_fn::{date_part, floor};
7use std::sync::Arc;
8use vegafusion_common::arrow::datatypes::{DataType, TimeUnit};
9use vegafusion_common::datafusion_common::DFSchema;
10use vegafusion_core::error::Result;
11
12pub fn make_local_datepart_transform(part: &str, tx: Option<fn(Expr) -> Expr>) -> TzTransformFn {
13    let part = part.to_string();
14    let local_datepart_transform =
15        move |tz_config: &RuntimeTzConfig, args: &[Expr], schema: &DFSchema| -> Result<Expr> {
16            let arg = args.first().unwrap().clone();
17            let arg = to_timestamp_col(arg, schema, &tz_config.default_input_tz.to_string())?;
18            let mut expr = date_part(
19                lit(part.clone()),
20                arg.try_cast_to(
21                    &DataType::Timestamp(
22                        TimeUnit::Millisecond,
23                        Some(tz_config.local_tz.to_string().into()),
24                    ),
25                    schema,
26                )?,
27            );
28
29            if let Some(tx) = tx {
30                expr = tx(expr)
31            }
32
33            Ok(expr)
34        };
35    Arc::new(local_datepart_transform)
36}
37
38pub fn make_utc_datepart_transform(part: &str, tx: Option<fn(Expr) -> Expr>) -> TzTransformFn {
39    let part = part.to_string();
40    let utc_datepart_transform =
41        move |tz_config: &RuntimeTzConfig, args: &[Expr], schema: &DFSchema| -> Result<Expr> {
42            let arg = to_timestamp_col(
43                args.first().unwrap().clone(),
44                schema,
45                &tz_config.default_input_tz.to_string(),
46            )?;
47            let mut expr = date_part(
48                lit(part.clone()),
49                arg.try_cast_to(
50                    &DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
51                    schema,
52                )?,
53            );
54
55            if let Some(tx) = tx {
56                expr = tx(expr)
57            }
58
59            Ok(expr)
60        };
61    Arc::new(utc_datepart_transform)
62}
63
64lazy_static! {
65    // Local Transforms
66    pub static ref YEAR_TRANSFORM: TzTransformFn =
67        make_local_datepart_transform("year", None);
68    pub static ref QUARTER_TRANSFORM: TzTransformFn =
69        make_local_datepart_transform("quarter", None);
70
71    // Months are zero-based in Vega
72    pub static ref MONTH_TRANSFORM: TzTransformFn =
73        make_local_datepart_transform(
74            "month", Some(|expr| expr - lit(1.0))
75        );
76
77    pub static ref DAYOFYEAR_TRANSFORM: TzTransformFn =
78        make_local_datepart_transform("doy", None);
79    pub static ref DATE_TRANSFORM: TzTransformFn =
80        make_local_datepart_transform("day", None);
81    pub static ref DAY_TRANSFORM: TzTransformFn =
82        make_local_datepart_transform("dow", None);
83    pub static ref HOUR_TRANSFORM: TzTransformFn =
84        make_local_datepart_transform("hour", None);
85    pub static ref MINUTE_TRANSFORM: TzTransformFn =
86        make_local_datepart_transform("minute", None);
87    pub static ref SECOND_TRANSFORM: TzTransformFn =
88        make_local_datepart_transform(
89            "second", Some(floor)
90        );
91    pub static ref MILLISECOND_TRANSFORM: TzTransformFn =
92        make_local_datepart_transform(
93            "millisecond",  Some(|expr| expr % lit(1000.0))
94        );
95
96    // UTC Transforms
97    pub static ref UTCYEAR_TRANSFORM: TzTransformFn =
98        make_utc_datepart_transform("year", None);
99    pub static ref UTCQUARTER_TRANSFORM: TzTransformFn =
100        make_utc_datepart_transform("quarter", None);
101
102    // Months are zero-based in Vega
103    pub static ref UTCMONTH_TRANSFORM: TzTransformFn =
104        make_utc_datepart_transform(
105            "month", Some(|expr| expr - lit(1.0))
106        );
107
108    pub static ref UTCDAYOFYEAR_TRANSFORM: TzTransformFn =
109        make_utc_datepart_transform("doy", None);
110    pub static ref UTCDATE_TRANSFORM: TzTransformFn =
111        make_utc_datepart_transform("day", None);
112    pub static ref UTCDAY_TRANSFORM: TzTransformFn =
113        make_utc_datepart_transform("dow", None);
114    pub static ref UTCHOUR_TRANSFORM: TzTransformFn =
115        make_utc_datepart_transform("hour", None);
116    pub static ref UTCMINUTE_TRANSFORM: TzTransformFn =
117        make_utc_datepart_transform("minute", None);
118    pub static ref UTCSECOND_TRANSFORM: TzTransformFn =
119        make_utc_datepart_transform(
120            "second", Some(floor)
121        );
122    pub static ref UTCMILLISECOND_TRANSFORM: TzTransformFn =
123        make_utc_datepart_transform(
124            "millisecond", Some(|expr| expr % lit(1000.0))
125        );
126}
vegafusion_runtime/expression/compiler/builtin_functions/date_time/date_parts.rs

vegafusion_runtime/expression/compiler/builtin_functions/date_time/
date_parts.rs
