Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore(query): add window function ast parser #10430

Merged
merged 6 commits into from
Mar 9, 2023
Merged

chore(query): add window function ast parser #10430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
50fd483
chore(query): add window parser
sundy-li Mar 7, 2023
20e6ac7
Merge branch 'main' into parser-window
sundy-li Mar 7, 2023
60651e0
chore(query): add window parser
sundy-li Mar 8, 2023
9465edc
chore(query): add window expression parser
sundy-li Mar 8, 2023
ce0f2a5
chore(query): add window expression parser
sundy-li Mar 8, 2023
866fa6b
chore(query): add window expression parser
sundy-li Mar 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 107 additions & 1 deletion src/query/ast/src/ast/expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use common_io::display_decimal_128;
use common_io::display_decimal_256;
use ethnum::i256;

use super::OrderByExpr;
use crate::ast::write_comma_separated_list;
use crate::ast::write_period_separated_list;
use crate::ast::Identifier;
Expand Down Expand Up @@ -146,14 +147,15 @@ pub enum Expr {
CountAll { span: Span },
/// `(foo, bar)`
Tuple { span: Span, exprs: Vec<Expr> },
/// Scalar function call
/// Scalar/Agg/Window function call
FunctionCall {
span: Span,
/// Set to true if the function is aggregate function with `DISTINCT`, like `COUNT(DISTINCT a)`
distinct: bool,
name: Identifier,
args: Vec<Expr>,
params: Vec<Literal>,
window: Option<WindowSpec>,
},
/// `CASE ... WHEN ... ELSE ...` expression
Case {
Expand Down Expand Up @@ -456,6 +458,38 @@ pub enum TrimWhere {
Trailing,
}

#[derive(Debug, Clone, PartialEq)]
pub struct WindowSpec {
pub partition_by: Vec<Expr>,
pub order_by: Vec<OrderByExpr>,
pub window_frame: Option<WindowFrame>,
}

/// `RANGE UNBOUNDED PRECEDING` or `ROWS BETWEEN 5 PRECEDING AND CURRENT ROW`.
#[derive(Debug, Clone, PartialEq)]
pub struct WindowFrame {
pub units: WindowFrameUnits,
pub start_bound: WindowFrameBound,
pub end_bound: WindowFrameBound,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WindowFrameUnits {
Rows,
Range,
}

/// Specifies [WindowFrame]'s `start_bound` and `end_bound`
#[derive(Debug, Clone, PartialEq)]
pub enum WindowFrameBound {
/// `CURRENT ROW`
CurrentRow,
/// `<N> PRECEDING` or `UNBOUNDED PRECEDING`
Preceding(Option<Box<Expr>>),
/// `<N> FOLLOWING` or `UNBOUNDED FOLLOWING`.
Following(Option<Box<Expr>>),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BinaryOperator {
Plus,
Expand Down Expand Up @@ -833,6 +867,73 @@ impl Display for Literal {
}
}

impl Display for WindowSpec {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let mut first = true;
if !self.partition_by.is_empty() {
first = false;
write!(f, "PARTITION BY ")?;
for (i, p) in self.partition_by.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{p}")?;
}
}

if !self.order_by.is_empty() {
if !first {
write!(f, " ")?;
}
first = false;
write!(f, "ORDER BY ")?;
for (i, o) in self.order_by.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{o}")?;
}
}

if let Some(frame) = &self.window_frame {
if !first {
write!(f, " ")?;
}
match frame.units {
WindowFrameUnits::Rows => {
write!(f, "ROWS")?;
}
WindowFrameUnits::Range => {
write!(f, "RANGE")?;
}
}
match (&frame.start_bound, &frame.end_bound) {
(WindowFrameBound::CurrentRow, WindowFrameBound::CurrentRow) => {
write!(f, " CURRENT ROW")?
}
_ => {
let format_frame = |frame: &WindowFrameBound| -> String {
match frame {
WindowFrameBound::CurrentRow => "CURRENT ROW".to_string(),
WindowFrameBound::Preceding(None) => "UNBOUNDED PRECEDING".to_string(),
WindowFrameBound::Following(None) => "UNBOUNDED FOLLOWING".to_string(),
WindowFrameBound::Preceding(Some(n)) => format!("{} PRECEDING", n),
WindowFrameBound::Following(Some(n)) => format!("{} FOLLOWING", n),
}
};
write!(
f,
" BETWEEN {} AND {}",
format_frame(&frame.start_bound),
format_frame(&frame.end_bound)
)?
}
}
}
Ok(())
}
}

impl Display for Expr {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Expand Down Expand Up @@ -981,6 +1082,7 @@ impl Display for Expr {
name,
args,
params,
window,
..
} => {
write!(f, "{name}")?;
Expand All @@ -995,6 +1097,10 @@ impl Display for Expr {
}
write_comma_separated_list(f, args)?;
write!(f, ")")?;

if let Some(window) = window {
write!(f, " OVER ({window})")?;
}
}
Expr::Case {
operand,
Expand Down
1 change: 1 addition & 0 deletions src/query/ast/src/ast/format/ast_format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,7 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
name: &'ast Identifier,
args: &'ast [Expr],
_params: &'ast [Literal],
_over: &'ast Option<WindowSpec>,
) {
let mut children = Vec::with_capacity(args.len());
for arg in args.iter() {
Expand Down
94 changes: 85 additions & 9 deletions src/query/ast/src/parser/expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -284,6 +284,7 @@ pub enum ExprElement {
distinct: bool,
name: Identifier,
args: Vec<Expr>,
window: Option<WindowSpec>,
params: Vec<Literal>,
},
/// `CASE ... WHEN ... ELSE ...` expression
Expand Down Expand Up @@ -476,12 +477,14 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
name,
args,
params,
window,
} => Expr::FunctionCall {
span: transform_span(elem.span.0),
distinct,
name,
args,
params,
window,
},
ExprElement::Case {
operand,
Expand Down Expand Up @@ -827,34 +830,86 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
}
},
);

let window_frame_between = alt((
map(
rule! { BETWEEN ~ #window_frame_bound ~ AND ~ #window_frame_bound },
|(_, s, _, e)| (s, e),
),
map(rule! {#window_frame_bound}, |s| {
(s, WindowFrameBound::Following(None))
}),
));

let window_spec = map(
rule! {
(PARTITION ~ ^BY ~ #comma_separated_list1(subexpr(0)))?
~ ( ORDER ~ ^BY ~ ^#comma_separated_list1(order_by_expr) )?
~ ((ROWS | RANGE) ~ #window_frame_between)?
},
|(opt_partition, opt_order, between)| WindowSpec {
partition_by: opt_partition.map(|x| x.2).unwrap_or_default(),
order_by: opt_order.map(|x| x.2).unwrap_or_default(),
window_frame: between.map(|x| {
let unit = match x.0.kind {
ROWS => WindowFrameUnits::Rows,
RANGE => WindowFrameUnits::Range,
_ => unreachable!(),
};
let bw = x.1;
WindowFrame {
units: unit,
start_bound: bw.0,
end_bound: bw.1,
}
}),
},
);

let function_call = map(
rule! {
#function_name
~ "("
~ DISTINCT?
~ #comma_separated_list0(subexpr(0))?
~ ")"
~ "(" ~ DISTINCT? ~ #comma_separated_list0(subexpr(0))? ~ ")"
},
|(name, _, opt_distinct, opt_args, _)| ExprElement::FunctionCall {
distinct: opt_distinct.is_some(),
name,
args: opt_args.unwrap_or_default(),
params: vec![],
window: None,
},
);
let function_call_with_param = map(

let function_call_with_window = map(
rule! {
#function_name
~ "(" ~ #comma_separated_list1(literal) ~ ")"
~ "(" ~ DISTINCT? ~ #comma_separated_list0(subexpr(0))? ~ ")"
~ (OVER ~ "(" ~ #window_spec ~ ")")
},
|(name, _, params, _, _, opt_distinct, opt_args, _)| ExprElement::FunctionCall {
|(name, _, opt_distinct, opt_args, _, window)| ExprElement::FunctionCall {
distinct: opt_distinct.is_some(),
name,
args: opt_args.unwrap_or_default(),
params,
params: vec![],
window: Some(window.2),
},
);

let function_call_with_params = map(
rule! {
#function_name
~ ("(" ~ #comma_separated_list1(literal) ~ ")")?
~ "(" ~ DISTINCT? ~ #comma_separated_list0(subexpr(0))? ~ ")"
},
|(name, params, _, opt_distinct, opt_args, _)| ExprElement::FunctionCall {
distinct: opt_distinct.is_some(),
name,
args: opt_args.unwrap_or_default(),
params: params.map(|x| x.1).unwrap_or_default(),
window: None,
},
);

let case = map(
rule! {
CASE ~ #subexpr(0)?
Expand Down Expand Up @@ -1036,7 +1091,8 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
| #trim_from : "`TRIM([(BOTH | LEADEING | TRAILING) ... FROM ...)`"
| #is_distinct_from: "`... IS [NOT] DISTINCT FROM ...`"
| #count_all : "COUNT(*)"
| #function_call_with_param : "<function>"
| #function_call_with_window : "<function>"
| #function_call_with_params : "<function>"
| #function_call : "<function>"
| #case : "`CASE ... END`"
| #subquery : "`(SELECT ...)`"
Expand All @@ -1052,6 +1108,26 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
Ok((rest, WithSpan { span, elem }))
}

pub fn window_frame_bound(i: Input) -> IResult<WindowFrameBound> {
alt((
value(WindowFrameBound::CurrentRow, rule! { CURRENT ~ ROW }),
map(rule! { #subexpr(0) ~ PRECEDING }, |(expr, _)| {
WindowFrameBound::Preceding(Some(Box::new(expr)))
}),
value(
WindowFrameBound::Preceding(None),
rule! { UNBOUNDED ~ PRECEDING },
),
map(rule! { #subexpr(0) ~ FOLLOWING }, |(expr, _)| {
WindowFrameBound::Following(Some(Box::new(expr)))
}),
value(
WindowFrameBound::Following(None),
rule! { UNBOUNDED ~ FOLLOWING },
),
))(i)
}

pub fn unary_op(i: Input) -> IResult<UnaryOperator> {
// Plus and Minus are parsed as binary op at first.
value(UnaryOperator::Not, rule! { NOT })(i)
Expand Down
Loading