use std::borrow::Cow;
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use ts_rs::TS;
use crate::config::{GroupRollupMode, SplitRollupMode};
use crate::proto::get_features_resp::{AggregateArgs, AggregateOptions, ColumnTypeOptions};
use crate::proto::{ColumnType, GetFeaturesResp, WindowAggregateArgs};
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize, TS)]
pub struct Features<'a> {
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub group_by: bool,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub group_rollup_mode: Vec<GroupRollupMode>,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub split_by: bool,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub split_rollup_mode: Vec<SplitRollupMode>,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub filter_ops: IndexMap<ColumnType, Vec<Cow<'a, str>>>,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub aggregates: IndexMap<ColumnType, Vec<AggSpec<'a>>>,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub sort: bool,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub expressions: bool,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub window_aggregates: IndexMap<ColumnType, Vec<WindowAggSpec<'a>>>,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub on_update: bool,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub unordered: bool,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize, TS)]
pub struct WindowAggSpec<'a> {
pub name: Cow<'a, str>,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub frames: Vec<Cow<'a, str>>,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub offset: bool,
#[serde(default)]
#[ts(optional, as = "Option<_>")]
pub alpha: bool,
#[serde(default)]
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub result_type: Option<ColumnType>,
}
impl<'a> From<&'a str> for WindowAggSpec<'a> {
fn from(name: &'a str) -> Self {
WindowAggSpec {
name: Cow::Borrowed(name),
frames: vec![],
offset: false,
alpha: false,
result_type: None,
}
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize, TS)]
#[serde(untagged)]
pub enum AggSpec<'a> {
Single(Cow<'a, str>),
Multiple(Cow<'a, str>, Vec<ColumnType>),
}
impl<'a> From<Features<'a>> for GetFeaturesResp {
fn from(value: Features<'a>) -> GetFeaturesResp {
GetFeaturesResp {
group_by: value.group_by,
group_rollup_mode: value
.group_rollup_mode
.iter()
.map(|x| crate::proto::GroupRollupMode::from(*x) as i32)
.collect(),
split_rollup_mode: value
.split_rollup_mode
.iter()
.map(|x| crate::proto::SplitRollupMode::from(*x) as i32)
.collect(),
split_by: value.split_by,
expressions: value.expressions,
on_update: value.on_update,
sort: value.sort,
unordered: value.unordered,
window_aggregates: value
.window_aggregates
.iter()
.map(|(ty, aggs)| {
(
*ty as u32,
crate::proto::get_features_resp::WindowAggregateOptions {
options: aggs
.iter()
.map(|x| WindowAggregateArgs {
name: x.name.to_string(),
frames: x.frames.iter().map(|f| f.to_string()).collect(),
offset: x.offset,
alpha: x.alpha,
result_type: x.result_type.map(|t| t as i32),
})
.collect(),
},
)
})
.collect(),
aggregates: value
.aggregates
.iter()
.map(|(dtype, aggs)| {
(*dtype as u32, AggregateOptions {
aggregates: aggs
.iter()
.map(|agg| match agg {
AggSpec::Single(cow) => AggregateArgs {
name: cow.to_string(),
args: vec![],
},
AggSpec::Multiple(cow, column_types) => AggregateArgs {
name: cow.to_string(),
args: column_types.iter().map(|x| *x as i32).collect(),
},
})
.collect(),
})
})
.collect(),
filter_ops: value
.filter_ops
.iter()
.map(|(ty, options)| {
(*ty as u32, ColumnTypeOptions {
options: options.iter().map(|x| (*x).to_string()).collect(),
})
})
.collect(),
}
}
}