Skip to content

Commit

Permalink
feat: support double quoted literal strings for dialects(such as mysq…
Browse files Browse the repository at this point in the history
…l,bigquery,spark) (#3056)

* support double quoted string.

* add test.

* add case sensitie test case.

* fix naming error.
  • Loading branch information
Rachelint authored Aug 8, 2022
1 parent c95bb74 commit 85d5363
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 5 deletions.
36 changes: 36 additions & 0 deletions datafusion/core/tests/sql/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1210,3 +1210,39 @@ async fn unprojected_filter() {
];
assert_batches_sorted_eq!(expected, &results);
}

#[tokio::test]
async fn case_sensitive_in_default_dialect() {
let int32_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
let schema = Schema::new(vec![Field::new("INT32", DataType::Int32, false)]);
let batch =
RecordBatch::try_new(Arc::new(schema), vec![Arc::new(int32_array)]).unwrap();

let ctx = SessionContext::new();
let table = MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap();
ctx.register_table("t", Arc::new(table)).unwrap();

{
let sql = "select \"int32\" from t";
let plan = ctx.create_logical_plan(sql);
assert!(plan.is_err());
}

{
let sql = "select \"INT32\" from t";
let actual = execute_to_batches(&ctx, sql).await;

let expected = vec![
"+-------+",
"| INT32 |",
"+-------+",
"| 1 |",
"| 2 |",
"| 3 |",
"| 4 |",
"| 5 |",
"+-------+",
];
assert_batches_eq!(expected, &actual);
}
}
43 changes: 38 additions & 5 deletions datafusion/sql/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,9 @@ fn plan_key(key: SQLExpr) -> Result<ScalarValue> {
SQLExpr::Value(Value::Number(s, _)) => {
ScalarValue::Int64(Some(s.parse().unwrap()))
}
SQLExpr::Value(Value::SingleQuotedString(s)) => ScalarValue::Utf8(Some(s)),
SQLExpr::Value(Value::SingleQuotedString(s) | Value::DoubleQuotedString(s)) => {
ScalarValue::Utf8(Some(s))
}
_ => {
return Err(DataFusionError::SQL(ParserError(format!(
"Unsuported index key expression: {:?}",
Expand Down Expand Up @@ -1596,7 +1598,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
row.into_iter()
.map(|v| match v {
SQLExpr::Value(Value::Number(n, _)) => parse_sql_number(&n),
SQLExpr::Value(Value::SingleQuotedString(s)) => Ok(lit(s)),
SQLExpr::Value(
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
) => Ok(lit(s)),
SQLExpr::Value(Value::Null) => {
Ok(Expr::Literal(ScalarValue::Null))
}
Expand Down Expand Up @@ -1638,7 +1642,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
) -> Result<Expr> {
match sql {
SQLExpr::Value(Value::Number(n, _)) => parse_sql_number(&n),
SQLExpr::Value(Value::SingleQuotedString(ref s)) => Ok(lit(s.clone())),
SQLExpr::Value(Value::SingleQuotedString(ref s) | Value::DoubleQuotedString(ref s)) => Ok(lit(s.clone())),
SQLExpr::Value(Value::Boolean(n)) => Ok(lit(n)),
SQLExpr::Value(Value::Null) => Ok(Expr::Literal(ScalarValue::Null)),
SQLExpr::Extract { field, expr } => Ok(Expr::ScalarFunction {
Expand Down Expand Up @@ -2219,7 +2223,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {

// Only handle string exprs for now
let value = match value {
SQLExpr::Value(Value::SingleQuotedString(s)) => s,
SQLExpr::Value(
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
) => s,
_ => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported interval argument. Expected string literal, got: {:?}",
Expand Down Expand Up @@ -2595,6 +2601,7 @@ fn parse_sql_number(n: &str) -> Result<Expr> {
mod tests {
use super::*;
use crate::assert_contains;
use sqlparser::dialect::{Dialect, GenericDialect, MySqlDialect};
use std::any::Any;

#[test]
Expand Down Expand Up @@ -4371,8 +4378,16 @@ mod tests {
}

fn logical_plan(sql: &str) -> Result<LogicalPlan> {
let dialect = &GenericDialect {};
logical_plan_with_dialect(sql, dialect)
}

fn logical_plan_with_dialect(
sql: &str,
dialect: &dyn Dialect,
) -> Result<LogicalPlan> {
let planner = SqlToRel::new(&MockContextProvider {});
let result = DFParser::parse_sql(sql);
let result = DFParser::parse_sql_with_dialect(sql, dialect);
let mut ast = result?;
planner.statement_to_plan(ast.pop_front().unwrap())
}
Expand Down Expand Up @@ -4840,6 +4855,24 @@ mod tests {
quick_test(sql, expected);
}

#[test]
fn test_double_quoted_literal_string() {
// Assert double quoted literal string is parsed correctly like single quoted one in specific dialect.
let dialect = &MySqlDialect {};
let single_quoted_res = format!(
"{:?}",
logical_plan_with_dialect("SELECT '1'", dialect).unwrap()
);
let double_quoted_res = format!(
"{:?}",
logical_plan_with_dialect("SELECT \"1\"", dialect).unwrap()
);
assert_eq!(single_quoted_res, double_quoted_res);

// It should return error in other dialect.
assert!(logical_plan("SELECT \"1\"").is_err());
}

fn assert_field_not_found(err: DataFusionError, name: &str) {
match err {
DataFusionError::SchemaError { .. } => {
Expand Down

0 comments on commit 85d5363

Please sign in to comment.