Influxdata
/
influxdb
mirror of https://github.com/influxdata/influxdb.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

chore: move walk and the mutable variant to the parser crate (#6575) 

It is generally a useful API to be core to the InfluxQL parser crate.
pull/24376/head
Stuart Carnie 2023-01-13 08:06:06 +11:00 committed by GitHub
parent 97e85a24e3
commit 81ffb3edb5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 282 additions and 173 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
influxdb_influxql_parser/src/expression.rs
View File

@ -7,6 +7,8 @@ pub use conditional::*;
pub mod arithmetic;
/// Provides conditional expression parsing.
pub mod conditional;
/// Provides APIs to traverse an expression tree using closures.
pub mod walk;
#[cfg(test)]
mod test_util;

4
influxdb_influxql_parser/src/expression/conditional.rs
View File

@ -15,7 +15,7 @@ use nom::sequence::{delimited, preceded, tuple};
use std::fmt;
use std::fmt::{Display, Formatter, Write};
/// Represents on of the conditional operators supported by [`ConditionalExpression::Binary`].
/// Represents one of the conditional operators supported by [`ConditionalExpression::Binary`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConditionalOperator {
/// Represents the `=` operator.
@ -233,7 +233,7 @@ impl ArithmeticParsers for ConditionalExpression {
}
/// Parse an arithmetic expression used by conditional expressions.
fn arithmetic_expression(i: &str) -> ParseResult<&str, Expr> {
pub(crate) fn arithmetic_expression(i: &str) -> ParseResult<&str, Expr> {
arithmetic::<ConditionalExpression>(i)
}

7
influxdb_influxql_parser/src/expression/snapshots/influxdb_influxql_parser__expression__walk__test__walk_expr-2.snap Normal file
View File

@ -0,0 +1,7 @@
---
source: influxdb_influxql_parser/src/expression/walk.rs
expression: "walk_expr(\"now() + 1h\")"
---
0: Call { name: "now", args: [] }
1: Literal(Duration(Duration(3600000000000)))
2: Binary { lhs: Call { name: "now", args: [] }, op: Add, rhs: Literal(Duration(Duration(3600000000000))) }

7
influxdb_influxql_parser/src/expression/snapshots/influxdb_influxql_parser__expression__walk__test__walk_expr.snap Normal file
View File

@ -0,0 +1,7 @@
---
source: influxdb_influxql_parser/src/expression/walk.rs
expression: "walk_expr(\"5 + 6\")"
---
0: Literal(Unsigned(5))
1: Literal(Unsigned(6))
2: Binary { lhs: Literal(Unsigned(5)), op: Add, rhs: Literal(Unsigned(6)) }

7
influxdb_influxql_parser/src/expression/snapshots/influxdb_influxql_parser__expression__walk__test__walk_expr_mut-2.snap Normal file
View File

@ -0,0 +1,7 @@
---
source: influxdb_influxql_parser/src/expression/walk.rs
expression: "walk_expr_mut(\"now() + 1h\")"
---
0: Call { name: "now", args: [] }
1: Literal(Duration(Duration(3600000000000)))
2: Binary { lhs: Call { name: "now", args: [] }, op: Add, rhs: Literal(Duration(Duration(3600000000000))) }

7
influxdb_influxql_parser/src/expression/snapshots/influxdb_influxql_parser__expression__walk__test__walk_expr_mut.snap Normal file
View File

@ -0,0 +1,7 @@
---
source: influxdb_influxql_parser/src/expression/walk.rs
expression: "walk_expr_mut(\"5 + 6\")"
---
0: Literal(Unsigned(5))
1: Literal(Unsigned(6))
2: Binary { lhs: Literal(Unsigned(5)), op: Add, rhs: Literal(Unsigned(6)) }

11
influxdb_influxql_parser/src/expression/snapshots/influxdb_influxql_parser__expression__walk__test__walk_expression-2.snap Normal file
View File

@ -0,0 +1,11 @@
---
source: influxdb_influxql_parser/src/expression/walk.rs
expression: "walk_expression(\"time > now() + 1h\")"
---
0: Arithmetic(VarRef { name: Identifier("time"), data_type: None })
1: Conditional(Expr(VarRef { name: Identifier("time"), data_type: None }))
2: Arithmetic(Call { name: "now", args: [] })
3: Arithmetic(Literal(Duration(Duration(3600000000000))))
4: Arithmetic(Binary { lhs: Call { name: "now", args: [] }, op: Add, rhs: Literal(Duration(Duration(3600000000000))) })
5: Conditional(Expr(Binary { lhs: Call { name: "now", args: [] }, op: Add, rhs: Literal(Duration(Duration(3600000000000))) }))
6: Conditional(Binary { lhs: Expr(VarRef { name: Identifier("time"), data_type: None }), op: Gt, rhs: Expr(Binary { lhs: Call { name: "now", args: [] }, op: Add, rhs: Literal(Duration(Duration(3600000000000))) }) })

13
influxdb_influxql_parser/src/expression/snapshots/influxdb_influxql_parser__expression__walk__test__walk_expression.snap Normal file
View File

@ -0,0 +1,13 @@
---
source: influxdb_influxql_parser/src/expression/walk.rs
expression: "walk_expression(\"5 + 6 = 2 + 9\")"
---
0: Arithmetic(Literal(Unsigned(5)))
1: Arithmetic(Literal(Unsigned(6)))
2: Arithmetic(Binary { lhs: Literal(Unsigned(5)), op: Add, rhs: Literal(Unsigned(6)) })
3: Conditional(Expr(Binary { lhs: Literal(Unsigned(5)), op: Add, rhs: Literal(Unsigned(6)) }))
4: Arithmetic(Literal(Unsigned(2)))
5: Arithmetic(Literal(Unsigned(9)))
6: Arithmetic(Binary { lhs: Literal(Unsigned(2)), op: Add, rhs: Literal(Unsigned(9)) })
7: Conditional(Expr(Binary { lhs: Literal(Unsigned(2)), op: Add, rhs: Literal(Unsigned(9)) }))
8: Conditional(Binary { lhs: Expr(Binary { lhs: Literal(Unsigned(5)), op: Add, rhs: Literal(Unsigned(6)) }), op: Eq, rhs: Expr(Binary { lhs: Literal(Unsigned(2)), op: Add, rhs: Literal(Unsigned(9)) }) })

207
influxdb_influxql_parser/src/expression/walk.rs Normal file
View File

@ -0,0 +1,207 @@
use crate::expression::{ConditionalExpression, Expr};
/// Expression distinguishes InfluxQL [`ConditionalExpression`] or [`Expr`]
/// nodes when visiting a [`ConditionalExpression`] tree. See [`walk_expression`].
#[derive(Debug)]
pub enum Expression<'a> {
/// Specifies a conditional expression.
Conditional(&'a ConditionalExpression),
/// Specifies an arithmetic expression.
Arithmetic(&'a Expr),
}
/// ExpressionMut is the same as [`Expression`] with the exception that
/// it provides mutable access to the nodes of the tree.
#[derive(Debug)]
pub enum ExpressionMut<'a> {
/// Specifies a conditional expression.
Conditional(&'a mut ConditionalExpression),
/// Specifies an arithmetic expression.
Arithmetic(&'a mut Expr),
}
/// Perform a depth-first traversal of an expression tree.
pub fn walk_expression<B>(
node: &ConditionalExpression,
visit: &mut impl FnMut(Expression<'_>) -> std::ops::ControlFlow<B>,
) -> std::ops::ControlFlow<B> {
match node {
ConditionalExpression::Expr(n) => walk_expr(n, &mut |n| visit(Expression::Arithmetic(n)))?,
ConditionalExpression::Binary { lhs, rhs, .. } => {
walk_expression(lhs, visit)?;
walk_expression(rhs, visit)?;
}
ConditionalExpression::Grouped(n) => walk_expression(n, visit)?,
}
visit(Expression::Conditional(node))
}
/// Perform a depth-first traversal of a mutable arithmetic or conditional expression tree.
pub fn walk_expression_mut<B>(
node: &mut ConditionalExpression,
visit: &mut impl FnMut(ExpressionMut<'_>) -> std::ops::ControlFlow<B>,
) -> std::ops::ControlFlow<B> {
match node {
ConditionalExpression::Expr(n) => {
walk_expr_mut(n, &mut |n| visit(ExpressionMut::Arithmetic(n)))?
}
ConditionalExpression::Binary { lhs, rhs, .. } => {
walk_expression_mut(lhs, visit)?;
walk_expression_mut(rhs, visit)?;
}
ConditionalExpression::Grouped(n) => walk_expression_mut(n, visit)?,
}
visit(ExpressionMut::Conditional(node))
}
/// Perform a depth-first traversal of the arithmetic expression tree.
pub fn walk_expr<B>(
expr: &Expr,
visit: &mut impl FnMut(&Expr) -> std::ops::ControlFlow<B>,
) -> std::ops::ControlFlow<B> {
match expr {
Expr::Binary { lhs, rhs, .. } => {
walk_expr(lhs, visit)?;
walk_expr(rhs, visit)?;
}
Expr::UnaryOp(_, n) => walk_expr(n, visit)?,
Expr::Nested(n) => walk_expr(n, visit)?,
Expr::Call { args, .. } => {
args.iter().try_for_each(|n| walk_expr(n, visit))?;
}
Expr::VarRef { .. }
| Expr::BindParameter(_)
| Expr::Literal(_)
| Expr::Wildcard(_)
| Expr::Distinct(_) => {}
}
visit(expr)
}
/// Perform a depth-first traversal of a mutable arithmetic expression tree.
pub fn walk_expr_mut<B>(
expr: &mut Expr,
visit: &mut impl FnMut(&mut Expr) -> std::ops::ControlFlow<B>,
) -> std::ops::ControlFlow<B> {
match expr {
Expr::Binary { lhs, rhs, .. } => {
walk_expr_mut(lhs, visit)?;
walk_expr_mut(rhs, visit)?;
}
Expr::UnaryOp(_, n) => walk_expr_mut(n, visit)?,
Expr::Nested(n) => walk_expr_mut(n, visit)?,
Expr::Call { args, .. } => {
args.iter_mut().try_for_each(|n| walk_expr_mut(n, visit))?;
}
Expr::VarRef { .. }
| Expr::BindParameter(_)
| Expr::Literal(_)
| Expr::Wildcard(_)
| Expr::Distinct(_) => {}
}
visit(expr)
}
#[cfg(test)]
mod test {
use crate::expression::walk::{walk_expr_mut, walk_expression_mut, ExpressionMut};
use crate::expression::{
arithmetic_expression, conditional_expression, ConditionalExpression, ConditionalOperator,
Expr,
};
use crate::literal::Literal;
#[test]
fn test_walk_expression() {
fn walk_expression(s: &str) -> String {
let (_, ref expr) = conditional_expression(s).unwrap();
let mut calls = Vec::new();
let mut call_no = 0;
super::walk_expression::<()>(expr, &mut |n| {
calls.push(format!("{}: {:?}", call_no, n));
call_no += 1;
std::ops::ControlFlow::Continue(())
});
calls.join("\n")
}
insta::assert_display_snapshot!(walk_expression("5 + 6 = 2 + 9"));
insta::assert_display_snapshot!(walk_expression("time > now() + 1h"));
}
#[test]
fn test_walk_expression_mut_modify() {
let (_, ref mut expr) = conditional_expression("foo + bar + 5 =~ /str/").unwrap();
walk_expression_mut::<()>(expr, &mut |e| {
match e {
ExpressionMut::Arithmetic(n) => match n {
Expr::VarRef { name, .. } => *name = format!("c_{}", name).into(),
Expr::Literal(Literal::Unsigned(v)) => *v *= 10,
Expr::Literal(Literal::Regex(v)) => *v = format!("c_{}", v.0).into(),
_ => {}
},
ExpressionMut::Conditional(n) => {
if let ConditionalExpression::Binary { op, .. } = n {
*op = ConditionalOperator::NotEqRegex
}
}
}
std::ops::ControlFlow::Continue(())
});
assert_eq!(format!("{}", expr), "c_foo + c_bar + 50 !~ /c_str/")
}
#[test]
fn test_walk_expr() {
fn walk_expr(s: &str) -> String {
let (_, expr) = arithmetic_expression(s).unwrap();
let mut calls = Vec::new();
let mut call_no = 0;
super::walk_expr::<()>(&expr, &mut |n| {
calls.push(format!("{}: {:?}", call_no, n));
call_no += 1;
std::ops::ControlFlow::Continue(())
});
calls.join("\n")
}
insta::assert_display_snapshot!(walk_expr("5 + 6"));
insta::assert_display_snapshot!(walk_expr("now() + 1h"));
}
#[test]
fn test_walk_expr_mut() {
fn walk_expr_mut(s: &str) -> String {
let (_, mut expr) = arithmetic_expression(s).unwrap();
let mut calls = Vec::new();
let mut call_no = 0;
super::walk_expr_mut::<()>(&mut expr, &mut |n| {
calls.push(format!("{}: {:?}", call_no, n));
call_no += 1;
std::ops::ControlFlow::Continue(())
});
calls.join("\n")
}
insta::assert_display_snapshot!(walk_expr_mut("5 + 6"));
insta::assert_display_snapshot!(walk_expr_mut("now() + 1h"));
}
#[test]
fn test_walk_expr_mut_modify() {
let (_, mut expr) = arithmetic_expression("foo + bar + 5").unwrap();
walk_expr_mut::<()>(&mut expr, &mut |e| {
match e {
Expr::VarRef { name, .. } => *name = format!("c_{}", name).into(),
Expr::Literal(Literal::Unsigned(v)) => *v *= 10,
_ => {}
}
std::ops::ControlFlow::Continue(())
});
assert_eq!(format!("{}", expr), "c_foo + c_bar + 50")
}
}

134
iox_query/src/plan/influxql/rewriter.rs
View File

@ -5,6 +5,7 @@ use crate::plan::influxql::field::field_name;
use crate::plan::influxql::field_mapper::{field_and_dimensions, FieldTypeMap, TagSet};
use datafusion::common::{DataFusionError, Result};
use influxdb_influxql_parser::common::{MeasurementName, QualifiedMeasurementName};
use influxdb_influxql_parser::expression::walk::{walk_expr, walk_expr_mut};
use influxdb_influxql_parser::expression::{Expr, VarRefDataType, WildcardType};
use influxdb_influxql_parser::identifier::Identifier;
use influxdb_influxql_parser::literal::Literal;
@ -19,7 +20,7 @@ use predicate::rpc_predicate::QueryNamespaceMeta;
use query_functions::clean_non_meta_escapes;
use std::borrow::Borrow;
use std::collections::{HashMap, HashSet};
use std::ops::Deref;
use std::ops::{ControlFlow, Deref};
fn parse_regex(re: &Regex) -> Result<regex::Regex> {
let pattern = clean_non_meta_escapes(re.as_str());
@ -184,50 +185,6 @@ fn has_wildcards(stmt: &SelectStatement) -> (bool, bool) {
(res.0, res.1)
}
/// Perform a depth-first traversal of the expression tree.
fn walk_expr_mut(expr: &mut Expr, visit: &mut impl FnMut(&mut Expr) -> Result<()>) -> Result<()> {
match expr {
Expr::Binary { lhs, rhs, .. } => {
walk_expr_mut(lhs, visit)?;
walk_expr_mut(rhs, visit)?;
}
Expr::UnaryOp(_, expr) => walk_expr_mut(expr, visit)?,
Expr::Nested(expr) => walk_expr_mut(expr, visit)?,
Expr::Call { args, .. } => {
args.iter_mut().try_for_each(|n| walk_expr_mut(n, visit))?;
}
Expr::VarRef { .. }
| Expr::BindParameter(_)
| Expr::Literal(_)
| Expr::Wildcard(_)
| Expr::Distinct(_) => {}
}
visit(expr)
}
/// Perform a depth-first traversal of the expression tree.
pub(crate) fn walk_expr(expr: &Expr, visit: &mut impl FnMut(&Expr) -> Result<()>) -> Result<()> {
match expr {
Expr::Binary { lhs, rhs, .. } => {
walk_expr(lhs, visit)?;
walk_expr(rhs, visit)?;
}
Expr::UnaryOp(_, expr) => walk_expr(expr, visit)?,
Expr::Nested(expr) => walk_expr(expr, visit)?,
Expr::Call { args, .. } => {
args.iter().try_for_each(|n| walk_expr(n, visit))?;
}
Expr::VarRef { .. }
| Expr::BindParameter(_)
| Expr::Literal(_)
| Expr::Wildcard(_)
| Expr::Distinct(_) => {}
}
visit(expr)
}
/// Rewrite the projection list and GROUP BY of the specified `SELECT` statement.
///
/// Wildcards and regular expressions in the `SELECT` projection list and `GROUP BY` are expanded.
@ -248,18 +205,23 @@ fn rewrite_field_list(
// Attempt to rewrite all variable references in the fields with their types, if one
// hasn't been specified.
stmt.fields.iter_mut().try_for_each(|f| {
walk_expr_mut(&mut f.expr, &mut |e| {
if let ControlFlow::Break(e) = stmt.fields.iter_mut().try_for_each(|f| {
walk_expr_mut::<DataFusionError>(&mut f.expr, &mut |e| {
if matches!(e, Expr::VarRef { .. }) {
let new_type = evaluate_type(namespace, e.borrow(), &stmt.from)?;
let new_type = match evaluate_type(namespace, e.borrow(), &stmt.from) {
Err(e) => ControlFlow::Break(e)?,
Ok(v) => v,
};
if let Expr::VarRef { data_type, .. } = e {
*data_type = new_type;
}
}
Ok(())
ControlFlow::Continue(())
})
})?;
}) {
return Err(e);
}
let (has_field_wildcard, has_group_by_wildcard) = has_wildcards(stmt);
if (has_field_wildcard, has_group_by_wildcard) == (false, false) {
@ -421,17 +383,16 @@ fn rewrite_field_list(
}
Expr::Binary { .. } => {
let mut has_wildcard = false;
walk_expr(&f.expr, &mut |e| {
let has_wildcard = walk_expr(&f.expr, &mut |e| {
match e {
Expr::Wildcard(_) | Expr::Literal(Literal::Regex(_)) => {
has_wildcard = true
return ControlFlow::Break(())
}
_ => {}
}
Ok(())
})?;
ControlFlow::Continue(())
})
.is_break();
if has_wildcard {
return Err(DataFusionError::External(
@ -539,10 +500,8 @@ pub(crate) fn rewrite_statement(
#[cfg(test)]
mod test {
use crate::plan::influxql::rewriter::{has_wildcards, rewrite_statement, walk_expr_mut};
use crate::plan::influxql::test_utils::{get_first_field, MockNamespace};
use influxdb_influxql_parser::expression::Expr;
use influxdb_influxql_parser::literal::Literal;
use crate::plan::influxql::rewriter::{has_wildcards, rewrite_statement};
use crate::plan::influxql::test_utils::MockNamespace;
use influxdb_influxql_parser::parse_statements;
use influxdb_influxql_parser::select::SelectStatement;
use influxdb_influxql_parser::statement::Statement;
@ -851,59 +810,4 @@ mod test {
assert!(!res.0);
assert!(!res.1);
}
#[test]
fn test_walk_expr() {
fn walk_expr(s: &str) -> String {
let expr = get_first_field(format!("SELECT {} FROM f", s).as_str()).expr;
let mut calls = Vec::new();
let mut call_no = 0;
super::walk_expr(&expr, &mut |n| {
calls.push(format!("{}: {}", call_no, n));
call_no += 1;
Ok(())
})
.unwrap();
calls.join("\n")
}
insta::assert_display_snapshot!(walk_expr("5 + 6"));
insta::assert_display_snapshot!(walk_expr("count(5, foo + 7)"));
insta::assert_display_snapshot!(walk_expr("count(5, foo + 7) + sum(bar)"));
}
#[test]
fn test_walk_expr_mut() {
fn walk_expr_mut(s: &str) -> String {
let mut expr = get_first_field(format!("SELECT {} FROM f", s).as_str()).expr;
let mut calls = Vec::new();
let mut call_no = 0;
super::walk_expr_mut(&mut expr, &mut |n| {
calls.push(format!("{}: {}", call_no, n));
call_no += 1;
Ok(())
})
.unwrap();
calls.join("\n")
}
insta::assert_display_snapshot!(walk_expr_mut("5 + 6"));
insta::assert_display_snapshot!(walk_expr_mut("count(5, foo + 7)"));
insta::assert_display_snapshot!(walk_expr_mut("count(5, foo + 7) + sum(bar)"));
}
#[test]
fn test_walk_expr_mut_modify() {
let mut expr = get_first_field("SELECT foo + bar + 5 FROM f").expr;
walk_expr_mut(&mut expr, &mut |e| {
match e {
Expr::VarRef { name, .. } => *name = format!("c_{}", name).into(),
Expr::Literal(Literal::Unsigned(v)) => *v *= 10,
_ => {}
}
Ok(())
})
.unwrap();
assert_eq!(format!("{}", expr), "c_foo + c_bar + 50")
}
}

9
iox_query/src/plan/influxql/snapshots/iox_query__plan__influxql__rewriter__test__walk_expr-2.snap
View File

@ -1,9 +0,0 @@
---
source: iox_query/src/plan/influxql/rewriter.rs
expression: "walk_expr(\"count(5, foo + 7)\")"
---
0: 5
1: foo
2: 7
3: foo + 7
4: count(5, foo + 7)

12
iox_query/src/plan/influxql/snapshots/iox_query__plan__influxql__rewriter__test__walk_expr-3.snap
View File

@ -1,12 +0,0 @@
---
source: iox_query/src/plan/influxql/rewriter.rs
expression: "walk_expr(\"count(5, foo + 7) + sum(bar)\")"
---
0: 5
1: foo
2: 7
3: foo + 7
4: count(5, foo + 7)
5: bar
6: sum(bar)
7: count(5, foo + 7) + sum(bar)

7
iox_query/src/plan/influxql/snapshots/iox_query__plan__influxql__rewriter__test__walk_expr.snap
View File

@ -1,7 +0,0 @@
---
source: iox_query/src/plan/influxql/rewriter.rs
expression: "walk_expr(\"5 + 6\")"
---
0: 5
1: 6
2: 5 + 6

9
iox_query/src/plan/influxql/snapshots/iox_query__plan__influxql__rewriter__test__walk_expr_mut-2.snap
View File

@ -1,9 +0,0 @@
---
source: iox_query/src/plan/influxql/rewriter.rs
expression: "walk_expr_mut(\"count(5, foo + 7)\")"
---
0: 5
1: foo
2: 7
3: foo + 7
4: count(5, foo + 7)

12
iox_query/src/plan/influxql/snapshots/iox_query__plan__influxql__rewriter__test__walk_expr_mut-3.snap
View File

@ -1,12 +0,0 @@
---
source: iox_query/src/plan/influxql/rewriter.rs
expression: "walk_expr_mut(\"count(5, foo + 7) + sum(bar)\")"
---
0: 5
1: foo
2: 7
3: foo + 7
4: count(5, foo + 7)
5: bar
6: sum(bar)
7: count(5, foo + 7) + sum(bar)

7
iox_query/src/plan/influxql/snapshots/iox_query__plan__influxql__rewriter__test__walk_expr_mut.snap
View File

@ -1,7 +0,0 @@
---
source: iox_query/src/plan/influxql/rewriter.rs
expression: "walk_expr_mut(\"5 + 6\")"
---
0: 5
1: 6
2: 5 + 6